diff --git a/tests/test_IO.py b/tests/test_IO.py index ab8b40abed..ccb29e9e93 100644 --- a/tests/test_IO.py +++ b/tests/test_IO.py @@ -5,12 +5,15 @@ import numpy as np import os from time import time +import xarray as xr from tidy3d import * from tidy3d import __version__ +import tidy3d as td from .utils import SIM_FULL as SIM from .utils import SIM_MONITORS as SIM2 from .utils import clear_tmp +from .test_data_monitor import make_flux_data # Store an example of every minor release simulation to test updater in the future SIM_DIR = "tests/sims" @@ -198,3 +201,10 @@ def test_yaml(): sim.to_yaml(path1) sim1 = Simulation.from_yaml(path1) assert sim1 == sim + + +def test_to_json_data(): + """Test that all simulations in ``SIM_DIR`` can be updated to current version and loaded.""" + data = make_flux_data() + assert json.loads(data._json_string())["flux"] is not None + assert json.loads(data._json_string(include_data=False))["flux"] is None diff --git a/tests/test_data.py b/tests/test_data.py deleted file mode 100644 index 2bc56002ab..0000000000 --- a/tests/test_data.py +++ /dev/null @@ -1,222 +0,0 @@ -"""Tests data.py""" - -import pytest -import numpy as np - -from tidy3d.components.data import * -import tidy3d as td -from tidy3d.log import DataError - -from .utils import clear_tmp - - -def test_scalar_field_data(): - f = np.linspace(1e14, 2e14, 1001) - x = np.linspace(-1, 1, 10) - y = np.linspace(-2, 2, 20) - z = np.linspace(0, 0, 1) - values = (1 + 1j) * np.random.random((len(x), len(y), len(z), len(f))) - data = ScalarFieldData(values=values, x=x, y=y, z=z, f=f) - _ = data.data - - -def test_scalar_field_time_data(): - t = np.linspace(0, 1e-12, 1001) - x = np.linspace(-1, 1, 10) - y = np.linspace(-2, 2, 20) - z = np.linspace(0, 0, 1) - values = np.random.random((len(x), len(y), len(z), len(t))) - data = ScalarFieldTimeData(values=values, x=x, y=y, z=z, t=t) - _ = data.data - - -def test_complex_scalar_field_time_data(): - t = np.linspace(0, 1e-12, 1001) - x = np.linspace(-1, 1, 10) - y = np.linspace(-2, 2, 20) - z = np.linspace(0, 0, 1) - values = (1 + 1j) * np.random.random((len(x), len(y), len(z), len(t))) - data = ScalarFieldTimeData(values=values, x=x, y=y, z=z, t=t) - _ = data.data - - -def test_scalar_permittivity_data(): - f = np.linspace(1e14, 2e14, 1001) - x = np.linspace(-1, 1, 10) - y = np.linspace(-2, 2, 20) - z = np.linspace(0, 0, 1) - values = (1 + 1j) * np.random.random((len(x), len(y), len(z), len(f))) - data = ScalarPermittivityData(values=values, x=x, y=y, z=z, f=f) - _ = data.data - - -def test_mode_amps_data(): - f = np.linspace(2e14, 3e14, 1001) - mode_index = np.arange(1, 3) - values = (1 + 1j) * np.random.random((2, len(f), len(mode_index))) - data = ModeAmpsData(values=values, direction=["+", "-"], mode_index=mode_index, f=f) - _ = data.data - - -def test_mode_index_data(): - f = np.linspace(2e14, 3e14, 1001) - mode_index = np.arange(1, 3) - values = (1 + 1j) * np.random.random((len(f), len(mode_index))) - data = ModeIndexData(values=values, f=f, mode_index=np.arange(1, 3)) - _ = data.n_eff - _ = data.k_eff - _ = data.n_complex - _ = data.data - - -def test_field_data(): - f = np.linspace(1e14, 2e14, 1001) - x = np.linspace(-1, 1, 10) - y = np.linspace(-2, 2, 20) - z = np.linspace(0, 0, 1) - values = (1 + 1j) * np.random.random((len(x), len(y), len(z), len(f))) - field = ScalarFieldData(values=values, x=x, y=y, z=z, f=f) - data = FieldData(data_dict={"Ex": field, "Ey": field}) - _ = data.data - - -def test_field_time_data(): - t = np.linspace(0, 1e-12, 1001) - x = np.linspace(-1, 1, 10) - y = np.linspace(-2, 2, 20) - z = np.linspace(0, 0, 1) - values = np.random.random((len(x), len(y), len(z), len(t))) - field = ScalarFieldTimeData(values=values, x=x, y=y, z=z, t=t) - data = FieldTimeData(data_dict={"Ex": field, "Ey": field}) - _ = data.data - - -def test_permittivity_data(): - f = np.linspace(1e14, 2e14, 1001) - x = np.linspace(-1, 1, 10) - y = np.linspace(-2, 2, 20) - z = np.linspace(0, 0, 1) - values = (1 + 1j) * np.random.random((len(x), len(y), len(z), len(f))) - eps = ScalarPermittivityData(values=values, x=x, y=y, z=z, f=f) - data = PermittivityData(data_dict={"eps_xx": eps, "eps_yy": eps, "eps_zz": eps}) - _ = data.eps_xx - _ = data.eps_yy - _ = data.eps_zz - _ = data.data - - -def test_flux_data(): - f = np.linspace(2e14, 3e14, 1001) - values = np.random.random((len(f),)) - data = FluxData(values=values, f=f) - _ = data.data - - -def test_mode_field_data(): - f = np.linspace(1e14, 2e14, 1001) - x = np.linspace(-1, 1, 10) - y = np.linspace(-2, 2, 20) - z = np.linspace(0, 0, 1) - mode_index = np.arange(0, 4) - values = (1 + 1j) * np.random.random((len(x), len(y), len(z), len(f), len(mode_index))) - field = ScalarModeFieldData(values=values, x=x, y=y, z=z, f=f, mode_index=mode_index) - data = ModeFieldData(data_dict={"Ex": field, "Ey": field}) - _ = data.sel_mode_index(1) - - -def make_sim_data(): - center = (0, 0, 0) - size = (2, 2, 2) - f0 = 1 - monitor = td.FieldMonitor(size=size, center=center, freqs=[f0], name="test") - - sim_size = (5, 5, 5) - sim = td.Simulation( - size=sim_size, - grid_spec=td.GridSpec.auto(wavelength=td.C_0 / f0), - monitors=[monitor], - run_time=1e-12, - sources=[ - td.PointDipole( - center=(0, 0, 0), - source_time=td.GaussianPulse(freq0=2e14, fwidth=2e13), - polarization="Ex", - ) - ], - ) - - def rand_data(): - return ScalarFieldData( - x=np.linspace(-1, 1, 10), - y=np.linspace(-1, 1, 10), - z=np.linspace(-1, 1, 10), - f=[f0], - values=np.random.random((10, 10, 10, 1)) * (1 + 1j), - ) - - fields = ["Ex", "Ey", "Ez", "Hx", "Hy", "Hz"] - data_dict = {field: rand_data() for field in fields} - field_data = FieldData(data_dict=data_dict) - - sim_data = SimulationData( - simulation=sim, monitor_data={"test": field_data}, log_string="field_decay :1e-4" - ) - return sim_data - - -@clear_tmp -def test_sim_data(): - - sim_data = make_sim_data() - f0 = float(sim_data["test"].Ex.f[0]) - - sim_data.plot_field("test", "Ex", val="real", x=0, freq=f0) - _ = sim_data.at_centers("test") - _ = sim_data.normalized - _ = sim_data.diverged - _ = sim_data.final_decay_value - _ = sim_data.log - _ = sim_data["test"].Ex - _ = sim_data["test"].Ey - _ = sim_data["test"].Ez - _ = sim_data["test"].Ex - _ = sim_data["test"].Ey - _ = sim_data["test"].Ez - _ = sim_data.normalize() - sim_data.to_file("tests/tmp/sim_data.hdf5") - sim_data2 = SimulationData.from_file("tests/tmp/sim_data.hdf5") - # assert sim_data == sim_data2 - - -def test_symmetries(): - f = np.linspace(1e14, 2e14, 1001) - x = np.linspace(-1, 1, 10) - y = np.linspace(-2, 2, 20) - z = np.linspace(0, 0, 1) - values = (1 + 1j) * np.random.random((len(x), len(y), len(z), len(f))) - field = ScalarFieldData(values=values, x=x, y=y, z=z, f=f) - data = FieldData( - data_dict={"Ex": field, "Ey": field}, symmetry=(1, -1, 0), symmetry_center=(0, 0, 0) - ) - _ = data.expand_syms - - -@clear_tmp -def test_data_io_raised(): - - sim_data = make_sim_data() - - with pytest.raises(DataError): - _ = sim_data._json_string() - - with pytest.raises(DataError): - sim_data.to_json("tmp/path") - - with pytest.raises(DataError): - sim_data.to_yaml("tmp/path") - - with pytest.raises(DataError): - SimulationData.from_json("tmp/path") - - with pytest.raises(DataError): - SimulationData.from_yaml("tmp/path") diff --git a/tidy3d/components/base.py b/tidy3d/components/base.py index 0a214cadfb..b193ef63f0 100644 --- a/tidy3d/components/base.py +++ b/tidy3d/components/base.py @@ -14,7 +14,7 @@ import xarray as xr from .types import ComplexNumber, Literal, TYPE_TAG_STR # , DataObject -from ..log import FileError +from ..log import FileError, log # default indentation (# spaces) in files INDENT = 4 @@ -61,6 +61,7 @@ class Config: # pylint: disable=too-few-public-methods json_encoders = { np.ndarray: lambda x: tuple(x.tolist()), complex: lambda x: ComplexNumber(real=x.real, imag=x.imag), + xr.DataArray: lambda x: x.to_dict(), } frozen = True allow_mutation = False @@ -118,23 +119,30 @@ def from_file(cls, fname: str, **parse_kwargs) -> Tidy3dBaseModel: raise FileError(f"File must be .json, .yaml, or .hdf5 type, given {fname}") - def to_file(self, fname: str) -> None: + def to_file(self, fname: str, include_data: bool = True) -> None: """Exports :class:`Tidy3dBaseModel` instance to .yaml or .json file Parameters ---------- fname : str Full path to the .yaml or .json file to save the :class:`Tidy3dBaseModel` to. + include_data : bool = True + Whether to include xarray data. Note: data is always included in .hdf5 file. Example ------- >>> simulation.to_file(fname='folder/sim.json') # doctest: +SKIP """ if ".json" in fname: - return self.to_json(fname=fname) + return self.to_json(fname=fname, include_data=include_data) if ".yaml" in fname: - return self.to_yaml(fname=fname) + return self.to_yaml(fname=fname, include_data=include_data) if ".hdf5" in fname: + if not include_data: + log.warning( + "`include_data` set to `False`." + "This will have no effect for `.hdf5` format and the data will still be written." + ) return self.to_hdf5(fname=fname) raise FileError(f"File must be .json, .yaml, or .hdf5 type, given {fname}") @@ -161,19 +169,21 @@ def from_json(cls, fname: str, **parse_file_kwargs) -> Tidy3dBaseModel: """ return cls.parse_file(fname, **parse_file_kwargs) - def to_json(self, fname: str) -> None: + def to_json(self, fname: str, include_data: bool = True) -> None: """Exports :class:`Tidy3dBaseModel` instance to .json file Parameters ---------- fname : str Full path to the .json file to save the :class:`Tidy3dBaseModel` to. + include_data : bool = True + Whether to include xarray data. Example ------- >>> simulation.to_json(fname='folder/sim.json') # doctest: +SKIP """ - json_string = self._json_string() + json_string = self._json_string(include_data=include_data) with open(fname, "w", encoding="utf-8") as file_handle: file_handle.write(json_string) @@ -202,19 +212,21 @@ def from_yaml(cls, fname: str, **parse_raw_kwargs) -> Tidy3dBaseModel: json_raw = json.dumps(json_dict, indent=INDENT) return cls.parse_raw(json_raw, **parse_raw_kwargs) - def to_yaml(self, fname: str) -> None: + def to_yaml(self, fname: str, include_data: bool = True) -> None: """Exports :class:`Tidy3dBaseModel` instance to .yaml file. Parameters ---------- fname : str Full path to the .yaml file to save the :class:`Tidy3dBaseModel` to. + include_data : bool = True + Whether to include xarray data. Example ------- >>> simulation.to_yaml(fname='folder/sim.yaml') # doctest: +SKIP """ - json_string = self._json_string() + json_string = self._json_string(include_data=include_data) json_dict = json.loads(json_string) with open(fname, "w+", encoding="utf-8") as file_handle: yaml.dump(json_dict, file_handle, indent=INDENT) @@ -458,13 +470,15 @@ def __ge__(self, other): """define >= for getting unique indices based on hash.""" return hash(self) >= hash(other) - def _json_string(self, include_unset: bool = True) -> str: + def _json_string(self, include_unset: bool = True, include_data: bool = True) -> str: """Returns string representation of a :class:`Tidy3dBaseModel`. Parameters ---------- include_unset : bool = True Whether to include default fields in json string. + include_data : bool = True + Whether to include ``xarray`` data. Returns ------- @@ -473,6 +487,11 @@ def _json_string(self, include_unset: bool = True) -> str: """ exclude_unset = not include_unset + # if not include_data, temporarily set the xr.DataArray encoder to return None + original_encoder = self.__config__.json_encoders[xr.DataArray] + if not include_data: + self.__config__.json_encoders[xr.DataArray] = lambda x: None + # put infinity and -infinity in quotes tmp_string = "<>" json_string = self.json(indent=INDENT, exclude_unset=exclude_unset) @@ -480,6 +499,9 @@ def _json_string(self, include_unset: bool = True) -> str: json_string = json_string.replace("Infinity", '"Infinity"') json_string = json_string.replace(tmp_string, '"-Infinity"') + # re-set the json encoder for data + self.__config__.json_encoders[xr.DataArray] = original_encoder + return json_string @classmethod