From db459f078791d84ca2e36bdd8e13ad53bcc485c4 Mon Sep 17 00:00:00 2001
From: momchil <momchil@flexcompute.com>
Date: Thu, 15 Aug 2024 13:49:55 +0200
Subject: [PATCH] Add Simulation.num_computational_grid_points property

---
 CHANGELOG.md                             |  3 ++-
 tests/test_components/test_simulation.py | 21 ++++++++++++++++
 tidy3d/components/simulation.py          | 32 +++++++++++++++++++++++-
 3 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 40a12afb8..935382dd9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Added value_and_grad function to the autograd plugin, importable via `from tidy3d.plugins.autograd import value_and_grad`. Supports differentiating functions with auxiliary data (`value_and_grad(f, has_aux=True)`).
-
+- `Simulation.num_computational_grid_points` property to examine the number of grid cells that compose the computational domain corresponding to the simulation. This can differ from `Simulation.num_cells` based on boundary conditions and symmetries.
+ 
 ### Fixed
 - `DataArray` interpolation failure due to incorrect ordering of coordinates when interpolating with autograd tracers.
 - Error in `CustomSourceTime` when evaluating at a list of times entirely outside of the range of the envelope definition times.
diff --git a/tests/test_components/test_simulation.py b/tests/test_components/test_simulation.py
index 84e72dc38..c4093d12c 100644
--- a/tests/test_components/test_simulation.py
+++ b/tests/test_components/test_simulation.py
@@ -108,6 +108,27 @@ def test_sim_init():
     sim.epsilon(m)
 
 
+def test_num_cells():
+    """Test num_cells and num_computational_grid_points."""
+
+    sim = td.Simulation(
+        size=(1, 1, 1),
+        run_time=1e-12,
+        grid_spec=td.GridSpec.uniform(dl=0.1),
+        sources=[
+            td.PointDipole(
+                center=(0, 0, 0),
+                polarization="Ex",
+                source_time=td.GaussianPulse(freq0=2e14, fwidth=1e14),
+            )
+        ],
+    )
+    assert sim.num_computational_grid_points > sim.num_cells  # due to extra pixels at boundaries
+
+    sim = sim.updated_copy(symmetry=(1, 0, 0))
+    assert sim.num_computational_grid_points < sim.num_cells  # due to symmetry
+
+
 def test_monitors_data_size():
     """make sure a simulation can be initialized"""
 
diff --git a/tidy3d/components/simulation.py b/tidy3d/components/simulation.py
index 4b3b1e82f..fe9e6ef62 100644
--- a/tidy3d/components/simulation.py
+++ b/tidy3d/components/simulation.py
@@ -4239,7 +4239,7 @@ def grid(self) -> Grid:
 
     @cached_property
     def num_cells(self) -> int:
-        """Number of cells in the simulation.
+        """Number of cells in the simulation grid.
 
         Returns
         -------
@@ -4249,6 +4249,36 @@ def num_cells(self) -> int:
 
         return np.prod(self.grid.num_cells, dtype=np.int64)
 
+    @property
+    def _num_computational_grid_points_dim(self):
+        """Number of cells in the computational domain for this simulation along each dimension."""
+        num_cells = self.grid.num_cells
+        num_cells_comp_domain = []
+        # symmetry overrides other boundaries so should be checked first
+        for sym, npts, boundary in zip(self.symmetry, num_cells, self.boundary_spec.to_list):
+            if sym != 0:
+                num_cells_comp_domain.append(npts // 2 + 2)
+            elif isinstance(boundary[0], Periodic):
+                num_cells_comp_domain.append(npts)
+            else:
+                num_cells_comp_domain.append(npts + 2)
+        return num_cells_comp_domain
+
+    @property
+    def num_computational_grid_points(self):
+        """Number of cells in the computational domain for this simulation. This is usually
+        different from ``num_cells`` due to the boundary conditions. Specifically, all boundary
+        conditions apart from ``Periodic`` require an extra pixel at the end of the simulation
+        domain. On the other hand, if a symmetry is present along a given dimension, only half of
+        the grid cells along that dimension will be in the computational domain.
+
+        Returns
+        -------
+        int
+            Number of yee cells in the computational domain corresponding to the simulation.
+        """
+        return np.prod(self._num_computational_grid_points_dim, dtype=np.int64)
+
     def get_refractive_indices(self, freq: float) -> list[float]:
         """List of refractive indices in the simulation at a given frequency."""