chore(tidy3d): FXC-4318-add-mypy-typedefs-in-rest-of-repo-except-big-component-files

Author: marcorudolphflex

pyproject.toml

@@ -341,17 +341,15 @@ python_files = "*.py"
 [tool.mypy]
 python_version = "3.10"
 files = [
-  "tidy3d/components/autograd",
-  "tidy3d/components/data",
-  "tidy3d/components/geometry",
-  "tidy3d/components/grid",
-  "tidy3d/components/mode",
-  "tidy3d/components/tcad",
-  "tidy3d/components/viz",
-  "tidy3d/config",
-  "tidy3d/material_library",
-  "tidy3d/plugins",
-  "tidy3d/web",
+  "tidy3d",
+]
+exclude = [
+  "tidy3d/components/simulation.py",
+  "tidy3d/components/validators.py",
+  "tidy3d/components/medium.py",
+  "tidy3d/components/boundary.py",
+  "tidy3d/components/base.py",
+  "tidy3d/components/monitor.py",
 ]
 ignore_missing_imports = true
 follow_imports = "skip"

tidy3d/__main__.py

@@ -9,7 +9,7 @@
 from tidy3d.web import Job
 
 
-def main(args) -> None:
+def main(args: list[str]) -> None:
     """Parse args and run the corresponding tidy3d simulaton."""
 
     parser = argparse.ArgumentParser(description="Tidy3D")

tidy3d/components/base_sim/data/sim_data.py

@@ -9,6 +9,7 @@
 import xarray as xr
 from pydantic import Field, field_validator, model_validator
 
+from tidy3d.compat import Self
 from tidy3d.components.base import Tidy3dBaseModel
 from tidy3d.components.base_sim.data.monitor_data import AbstractMonitorData
 from tidy3d.components.base_sim.simulation import AbstractSimulation
@@ -51,7 +52,7 @@ def monitor_data(self) -> dict[str, AbstractMonitorData]:
         return {monitor_data.monitor.name: monitor_data for monitor_data in self.data}
 
     @model_validator(mode="after")
-    def data_monitors_match_sim(self):
+    def data_monitors_match_sim(self) -> Self:
         """Ensure each :class:`AbstractMonitorData` in ``.data`` corresponds to a monitor in
         ``.simulation``.
         """
@@ -70,7 +71,9 @@ def data_monitors_match_sim(self):
 
     @field_validator("data")
     @classmethod
-    def validate_no_ambiguity(cls, val):
+    def validate_no_ambiguity(
+        cls, val: tuple[AbstractMonitorData, ...]
+    ) -> tuple[AbstractMonitorData, ...]:
         """Ensure all :class:`AbstractMonitorData` entries in ``.data`` correspond to different
         monitors in ``.simulation``.
         """

tidy3d/components/base_sim/simulation.py

@@ -137,7 +137,7 @@ class AbstractSimulation(Box, ABC):
 
     @model_validator(mode="before")
     @classmethod
-    def _update_simulation(cls, data):
+    def _update_simulation(cls, data: dict[str, Any]) -> dict[str, Any]:
         """Update the simulation if it is an earlier version."""
         # dummy upgrade of version number
         # this should be overriden by each simulation class if needed
@@ -161,7 +161,7 @@ def _update_simulation(cls, data):
     _warn_traced_size = _warn_unsupported_traced_argument("size")
 
     @model_validator(mode="after")
-    def _structures_not_at_edges(self):
+    def _structures_not_at_edges(self) -> Self:
         """Warn if any structures lie at the simulation boundaries."""
 
         if self.structures is None:
@@ -704,7 +704,7 @@ def from_scene(cls, scene: Scene, **kwargs: Any) -> AbstractSimulation:
             **kwargs,
         )
 
-    def plot_3d(self, width=800, height=800) -> None:
+    def plot_3d(self, width: float = 800, height: float = 800) -> None:
         """Render 3D plot of ``AbstractSimulation`` (in jupyter notebook only).
         Parameters
         ----------

tidy3d/components/bc_placement.py

@@ -47,7 +47,7 @@ class StructureStructureInterface(AbstractBCPlacement):
 
     @field_validator("structures")
     @classmethod
-    def unique_names(cls, val):
+    def unique_names(cls, val: tuple[str, str]) -> tuple[str, str]:
         """Error if the same structure is provided twice"""
         if val[0] == val[1]:
             raise SetupError(
@@ -71,7 +71,7 @@ class MediumMediumInterface(AbstractBCPlacement):
 
     @field_validator("mediums")
     @classmethod
-    def unique_names(cls, val):
+    def unique_names(cls, val: tuple[str, str]) -> tuple[str, str]:
         """Error if the same structure is provided twice"""
         if val[0] == val[1]:
             raise SetupError("The same medium is provided twice in 'MediumMediumInterface'.")

tidy3d/components/beam.py

@@ -7,6 +7,7 @@
 from typing import Literal, Optional, Union
 
 import autograd.numpy as np
+from numpy.typing import NDArray
 from pydantic import Field, PositiveFloat
 
 from tidy3d.constants import C_0, ETA_0, HERTZ, MICROMETER, RADIAN
@@ -131,7 +132,9 @@ def field_data(self) -> FieldData:
 
         return data_raw.updated_copy(**fields_norm)
 
-    def _field_data_on_grid(self, grid: Grid, background_n: np.ndarray, colocate=True) -> dict:
+    def _field_data_on_grid(
+        self, grid: Grid, background_n: NDArray, colocate: bool = True
+    ) -> dict[str, ScalarFieldDataArray]:
         """Compute the field data for each field component on a grid for the beam.
         A dictionary of the scalar field data arrays is returned, not yet packaged as ``FieldData``.
         """
@@ -165,14 +168,14 @@ def _field_data_on_grid(self, grid: Grid, background_n: np.ndarray, colocate=Tru
         return scalar_fields
 
     @abstractmethod
-    def scalar_field(self, points: np.ndarray, background_n: float) -> np.ndarray:
+    def scalar_field(self, points: NDArray, background_n: float) -> NDArray:
         """Scalar field corresponding to the analytic beam in coordinate system such that the
         propagation direction is z and the ``E``-field is entirely ``x``-polarized. The field is
         computed on an unstructured array ``points`` of shape ``(3, ...)``."""
 
     def analytic_beam_z_normal(
-        self, points: np.ndarray, background_n: float, field: Literal["E", "H"]
-    ) -> np.ndarray:
+        self, points: NDArray, background_n: float, field: Literal["E", "H"]
+    ) -> NDArray:
         """Analytic beam with all the beam parameters but assuming ``z`` as the normal axis."""
 
         # Add a frequency dimension to points
@@ -212,12 +215,12 @@ def analytic_beam_z_normal(
 
     def analytic_beam(
         self,
-        x: np.ndarray,
-        y: np.ndarray,
-        z: np.ndarray,
+        x: NDArray,
+        y: NDArray,
+        z: NDArray,
         background_n: float,
         field: Literal["E", "H"],
-    ) -> np.ndarray:
+    ) -> NDArray:
         """Sample the analytic beam fields on a cartesian grid of points in x, y, z."""
 
         # Make a meshgrid
@@ -241,15 +244,13 @@ def analytic_beam(
         # Reshape to (3, Nx, Ny, Nz, num_freqs)
         return np.reshape(field_vals, (3, Nx, Ny, Nz, len(self.freqs)))
 
-    def _rotate_points_z(self, points: np.ndarray, background_n: np.ndarray) -> np.ndarray:
+    def _rotate_points_z(self, points: NDArray, background_n: NDArray) -> NDArray:
         """Rotate points to new coordinates where z is the propagation axis."""
         points_prop_z = self.rotate_points(points, [0, 0, 1], -self.angle_phi)
         points_prop_z = self.rotate_points(points_prop_z, [0, 1, 0], -self.angle_theta)
         return points_prop_z
 
-    def _inverse_rotate_field_vals_z(
-        self, field_vals: np.ndarray, background_n: np.ndarray
-    ) -> np.ndarray:
+    def _inverse_rotate_field_vals_z(self, field_vals: NDArray, background_n: NDArray) -> NDArray:
         """Rotate field values from coordinates where z is the propagation axis to angled
         coordinates."""
         field_vals = self.rotate_points(field_vals, [0, 1, 0], self.angle_theta)
@@ -288,18 +289,18 @@ class PlaneWaveBeamProfile(BeamProfile):
     )
 
     @property
-    def _angle_theta_frequency(self):
+    def _angle_theta_frequency(self) -> float:
         if not self.angle_theta_frequency:
             return np.mean(self.freqs)
         return self.angle_theta_frequency
 
-    def in_plane_k(self, background_n: float):
+    def in_plane_k(self, background_n: float) -> list[float]:
         """In-plane wave vector. Only the real part is taken so the beam has no in-plane decay."""
         k0 = 2 * np.pi * self._angle_theta_frequency / C_0 * background_n
         k_in_plane = k0.real * np.sin(self.angle_theta)
         return [k_in_plane * np.cos(self.angle_phi), k_in_plane * np.sin(self.angle_phi)]
 
-    def scalar_field(self, points: np.ndarray, background_n: float) -> np.ndarray:
+    def scalar_field(self, points: NDArray, background_n: float) -> NDArray:
         """Scalar field for plane wave.
         Scalar field corresponding to the analytic beam in coordinate system such that the
         propagation direction is z and the ``E``-field is entirely ``x``-polarized. The field is
@@ -314,14 +315,14 @@ def scalar_field(self, points: np.ndarray, background_n: float) -> np.ndarray:
             kz *= np.cos(self.angle_theta)
         return np.exp(1j * points[2] * kz)
 
-    def _angle_theta_actual(self, background_n: np.ndarray) -> np.ndarray:
+    def _angle_theta_actual(self, background_n: NDArray) -> NDArray:
         """Compute the frequency-dependent actual propagation angle theta."""
         k0 = 2 * np.pi * np.array(self.freqs) / C_0 * background_n
         kx, ky = self.in_plane_k(background_n)
         k_perp = np.sqrt(kx**2 + ky**2)
         return np.real(np.arcsin(k_perp / k0)) * np.sign(self.angle_theta)
 
-    def _rotate_points_z(self, points: np.ndarray, background_n: np.ndarray) -> np.ndarray:
+    def _rotate_points_z(self, points: NDArray, background_n: NDArray) -> NDArray:
         """Rotate points to new coordinates where z is the propagation axis."""
         if self.as_fixed_angle_source:
             # For fixed-angle, we do not rotate the points
@@ -335,9 +336,7 @@ def _rotate_points_z(self, points: np.ndarray, background_n: np.ndarray) -> np.n
             return points
         return super()._rotate_points_z(points, background_n)
 
-    def _inverse_rotate_field_vals_z(
-        self, field_vals: np.ndarray, background_n: np.ndarray
-    ) -> np.ndarray:
+    def _inverse_rotate_field_vals_z(self, field_vals: NDArray, background_n: NDArray) -> NDArray:
         """Rotate field values from coordinates where z is the propagation axis to angled
         coordinates. Special handling is needed if fixed in-plane k wave."""
         if isinstance(self.angular_spec, FixedInPlaneKSpec):
@@ -378,9 +377,7 @@ class GaussianBeamProfile(BeamProfile):
         units=MICROMETER,
     )
 
-    def beam_params(
-        self, z: np.ndarray, k0: np.ndarray
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    def beam_params(self, z: NDArray, k0: NDArray) -> tuple[NDArray, NDArray, NDArray]:
         """Compute the parameters needed to evaluate a Gaussian beam at z.
 
         Parameters
@@ -402,7 +399,7 @@ def beam_params(
         psi_g = np.arctan((z + z_0) / z_r) - np.arctan(z_0 / z_r)
         return w_z, inv_r_z, psi_g
 
-    def scalar_field(self, points: np.ndarray, background_n: float) -> np.ndarray:
+    def scalar_field(self, points: NDArray, background_n: float) -> NDArray:
         """Scalar field for Gaussian beam.
         Scalar field corresponding to the analytic beam in coordinate system such that the
         propagation direction is z and the ``E``-field is entirely ``x``-polarized. The field is
@@ -446,9 +443,7 @@ class AstigmaticGaussianBeamProfile(BeamProfile):
         units=MICROMETER,
     )
 
-    def beam_params(
-        self, z: np.ndarray, k0: np.ndarray
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    def beam_params(self, z: NDArray, k0: NDArray) -> tuple[NDArray, NDArray, NDArray, NDArray]:
         """Compute the parameters needed to evaluate an astigmatic Gaussian beam at z.
 
         Parameters
@@ -475,7 +470,7 @@ def beam_params(
 
         return w_0, w_z, inv_r_z, psi_g
 
-    def scalar_field(self, points: np.ndarray, background_n: float) -> np.ndarray:
+    def scalar_field(self, points: NDArray, background_n: float) -> NDArray:
         """
         Scalar field for astigmatic Gaussian beam.
         Scalar field corresponding to the analytic beam in coordinate system such that the

tidy3d/components/data/dataset.py

@@ -3,13 +3,14 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Literal, Optional, Self, Union, get_args
+from typing import Any, Callable, Literal, Optional, Union, get_args
 
 import numpy as np
 import xarray as xr
 from numpy.typing import ArrayLike
 from pydantic import Field
 
+from tidy3d.compat import Self
 from tidy3d.components.base import Tidy3dBaseModel
 from tidy3d.components.types import Axis, FreqArray, xyz
 from tidy3d.constants import C_0, PICOSECOND_PER_NANOMETER_PER_KILOMETER, UnitScaling

tidy3d/components/data/monitor_data.py

@@ -7,14 +7,15 @@
 from abc import ABC
 from math import isclose
 from os import PathLike
-from typing import Any, Callable, Literal, Optional, Self, SupportsComplex, Union, get_args
+from typing import Any, Callable, Literal, Optional, SupportsComplex, Union, get_args
 
 import autograd.numpy as np
 import xarray as xr
 from numpy.typing import NDArray
 from pandas import DataFrame
 from pydantic import Field, model_validator
 
+from tidy3d.compat import Self
 from tidy3d.components.base import cached_property
 from tidy3d.components.base_sim.data.monitor_data import AbstractMonitorData
 from tidy3d.components.grid.grid import Coords, Grid
@@ -265,7 +266,7 @@ def symmetry_expanded(self) -> Self:
         return self.updated_copy(**self._symmetry_update_dict, deep=False, validate=False)
 
     @property
-    def symmetry_expanded_copy(self) -> AbstractFieldData:
+    def symmetry_expanded_copy(self) -> Self:
         """Create a copy of the :class:`.AbstractFieldData` with fields expanded based on symmetry.
 
         Returns

tidy3d/components/data/unstructured/base.py

@@ -5,7 +5,7 @@
 import numbers
 from abc import ABC, abstractmethod
 from os import PathLike
-from typing import TYPE_CHECKING, Any, Literal, Optional, Self, Union
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
 import numpy as np
 from numpy.typing import DTypeLike, NDArray
@@ -14,6 +14,7 @@
 from vtkmodules.vtkCommonCore import vtkPoints
 from xarray import DataArray as XrDataArray
 
+from tidy3d.compat import Self
 from tidy3d.components.base import cached_property
 from tidy3d.components.data.data_array import (
     DATA_ARRAY_MAP,

tidy3d/components/data/unstructured/triangular.py

@@ -2,13 +2,15 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, Optional, Self, Union
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
 import numpy as np
 from pydantic import Field, PositiveInt
 from xarray import DataArray
 from xarray import DataArray as XrDataArray
 
+from tidy3d.compat import Self
+
 try:
     from matplotlib import pyplot as plt
     from matplotlib.tri import Triangulation