diff --git a/tidy3d/components/monitor.py b/tidy3d/components/monitor.py index d432dd581..7a47f1dc0 100644 --- a/tidy3d/components/monitor.py +++ b/tidy3d/components/monitor.py @@ -1,10 +1,11 @@ """Objects that define how data is recorded from simulation.""" from abc import ABC, abstractmethod -from typing import List, Union +from typing import List, Union, Tuple import pydantic +import numpy as np -from .types import Literal, Ax, EMField, ArrayLike +from .types import Literal, Ax, EMField, ArrayLike, Array from .geometry import Box from .validators import assert_plane from .mode import ModeSpec @@ -49,15 +50,15 @@ def geometry(self): return Box(center=self.center, size=self.size) @abstractmethod - def storage_size(self, num_cells: int, num_steps: int) -> int: + def storage_size(self, num_cells: int, tmesh: Array) -> int: """Size of monitor storage given the number of points after discretization. Parameters ---------- num_cells : int Number of grid cells within the monitor after discretization by a :class:`Simulation`. - num_steps : int - Number of time steps in the discretized :class:`Simulation`. + tmesh : Array + The discretized time mesh of a :class:`Simulation`. Returns ------- @@ -116,6 +117,42 @@ def stop_greater_than_start(cls, val, values): raise SetupError("Monitor start time is greater than stop time.") return val + def time_inds(self, tmesh: Array) -> Tuple[int, int]: + """Compute the starting and stopping index of the monitor in a given discrete time mesh.""" + + tind_beg, tind_end = (0, 0) + + if tmesh.size == 0: + return (tind_beg, tind_end) + + # If monitor.stop is None, record until the end + t_stop = self.stop + if t_stop is None: + tind_end = int(tmesh.size) + t_stop = tmesh[-1] + else: + tend = np.nonzero(tmesh <= t_stop)[0] + if tend.size > 0: + tind_end = int(tend[-1] + 1) + + # Step to compare to in order to handle t_start = t_stop + if np.array(tmesh).size < 2: + dt = 1e-20 + else: + dt = tmesh[1] - tmesh[0] + + # If equal start and stopping time, record one time step + if np.abs(self.start - t_stop) < dt: + tind_beg = max(tind_end - 1, 0) + else: + tbeg = np.nonzero(tmesh[0:tind_end] >= self.start)[0] + if tbeg.size > 0: + tind_beg = tbeg[0] + else: + tind_beg = tind_end + + return (tind_beg, tind_end) + class AbstractFieldMonitor(Monitor, ABC): """:class:`Monitor` that records electromagnetic field data as a function of x,y,z.""" @@ -221,7 +258,7 @@ class FieldMonitor(AbstractFieldMonitor, FreqMonitor): _data_type: Literal["ScalarFieldData"] = pydantic.Field("ScalarFieldData") - def storage_size(self, num_cells: int, num_steps: int) -> int: + def storage_size(self, num_cells: int, tmesh: Array) -> int: # stores 1 complex number per grid cell, per frequency, per field return BYTES_COMPLEX * num_cells * len(self.freqs) * len(self.fields) @@ -243,9 +280,11 @@ class FieldTimeMonitor(AbstractFieldMonitor, TimeMonitor): _data_type: Literal["ScalarFieldTimeData"] = pydantic.Field("ScalarFieldTimeData") - def storage_size(self, num_cells: int, num_steps: int) -> int: + def storage_size(self, num_cells: int, tmesh: Array) -> int: # stores 1 real number per grid cell, per time step, per field - return BYTES_REAL * num_cells * num_steps * len(self.fields) + time_inds = self.time_inds(tmesh) + num_steps = time_inds[1] - time_inds[0] + return BYTES_REAL * num_steps * num_cells * len(self.fields) class FluxMonitor(AbstractFluxMonitor, FreqMonitor): @@ -262,7 +301,7 @@ class FluxMonitor(AbstractFluxMonitor, FreqMonitor): _data_type: Literal["FluxData"] = pydantic.Field("FluxData") - def storage_size(self, num_cells: int, num_steps: int) -> int: + def storage_size(self, num_cells: int, tmesh: Array) -> int: # stores 6 complex numbers per grid cell, per frequency return 6 * BYTES_REAL * num_cells * len(self.freqs) @@ -283,8 +322,10 @@ class FluxTimeMonitor(AbstractFluxMonitor, TimeMonitor): _data_type: Literal["FluxTimeData"] = pydantic.Field("FluxTimeData") - def storage_size(self, num_cells: int, num_steps: int) -> int: + def storage_size(self, num_cells: int, tmesh: Array) -> int: # stores 1 real number per time tep + time_inds = self.time_inds(tmesh) + num_steps = time_inds[1] - time_inds[0] return BYTES_REAL * num_steps @@ -310,7 +351,7 @@ class ModeMonitor(PlanarMonitor, FreqMonitor): _data_type: Literal["ModeData"] = pydantic.Field("ModeData") - def storage_size(self, num_cells: int, num_steps: int) -> int: + def storage_size(self, num_cells: int, tmesh: int) -> int: # stores 3 complex numbers per grid cell, per frequency, per mode. return 3 * BYTES_COMPLEX * num_cells * len(self.freqs) * self.mode_spec.num_modes diff --git a/tidy3d/components/simulation.py b/tidy3d/components/simulation.py index 5dcc5457d..ac8446eaa 100644 --- a/tidy3d/components/simulation.py +++ b/tidy3d/components/simulation.py @@ -461,13 +461,13 @@ def _validate_size(self) -> None: def _validate_monitor_size(self) -> None: """Ensures the monitors arent storing too much data before simulation is uploaded.""" - num_time_steps = self.num_time_steps + tmesh = self.tmesh total_size_bytes = 0 for monitor in self.monitors: monitor_grid = self.discretize(monitor) num_cells = np.prod(monitor_grid.num_cells) - monitor_size = monitor.storage_size(num_cells=num_cells, num_steps=num_time_steps) + monitor_size = monitor.storage_size(num_cells=num_cells, tmesh=tmesh) total_size_bytes += monitor_size