fix(tidy3d): lazy load scipy to reduce import time (#2543)

Author: frederikschubertflex

tidy3d/components/dispersion_fitter.py

@@ -5,7 +5,6 @@
 from typing import Optional
 
 import numpy as np
-import scipy
 from pydantic.v1 import Field, NonNegativeFloat, PositiveFloat, PositiveInt, validator
 from rich.progress import Progress
 
@@ -493,6 +492,7 @@ def real_weighted_matrix(self, matrix: ArrayComplex2D) -> ArrayFloat2D:
 
     def iterate_poles(self) -> FastFitterData:
         """Perform a single iteration of the pole-updating procedure."""
+        from scipy import optimize
 
         def compute_zeros(residues: ArrayComplex1D, d_tilde: float) -> ArrayComplex1D:
             """Compute the zeros from the residues."""
@@ -564,7 +564,7 @@ def compute_zeros(residues: ArrayComplex1D, d_tilde: float) -> ArrayComplex1D:
                 )
 
             # solve the least squares problem
-            x_vector = scipy.optimize.lsq_linear(a_matrix_real, b_vector_real).x
+            x_vector = optimize.lsq_linear(a_matrix_real, b_vector_real).x
 
             # unpack the solution
             residues = np.zeros(len(self.poles), dtype=complex)
@@ -594,6 +594,8 @@ def compute_zeros(residues: ArrayComplex1D, d_tilde: float) -> ArrayComplex1D:
 
     def fit_residues(self) -> FastFitterData:
         """Fit residues."""
+        from scipy import optimize
+
         # build the matrices
         if self.optimize_eps_inf:
             poly_len = 1
@@ -610,7 +612,7 @@ def fit_residues(self) -> FastFitterData:
         # solve the least squares problem
         bounds = (-np.inf * np.ones(a_matrix.shape[1]), np.inf * np.ones(a_matrix.shape[1]))
         bounds[0][-1] = 1  # eps_inf >= 1
-        x_vector = scipy.optimize.lsq_linear(a_matrix_real, b_vector_real).x
+        x_vector = optimize.lsq_linear(a_matrix_real, b_vector_real).x
 
         # unpack the solution
         residues = np.zeros(len(self.poles), dtype=complex)
@@ -650,6 +652,7 @@ def iterate_fit(self) -> FastFitterData:
 
     def iterate_passivity(self, passivity_omega: ArrayFloat1D) -> tuple[FastFitterData, int]:
         """Iterate passivity enforcement algorithm."""
+        from scipy import optimize
 
         size = len(self.real_poles) + 2 * len(self.complex_poles)
         constraint_matrix = np.imag(self.pole_matrix_omega(passivity_omega))
@@ -683,7 +686,7 @@ def jac(dx):
 
         x0 = np.zeros(size)
         err = np.amin(c_vector - constraint_matrix @ x0)
-        result = scipy.optimize.minimize(
+        result = optimize.minimize(
             loss, x0=x0, jac=jac, constraints=cons, method="SLSQP", options=opt
         )
         x_vector = result.x

tidy3d/components/geometry/base.py

@@ -1386,7 +1386,7 @@ def to_gds(
         import gdstk
 
         if not isinstance(cell, gdstk.Cell):
-            if "gdstk" in cell.__class__.__name__.lower():  # type: ignore[attr-defined]
+            if "gdstk" in cell.__class__.__name__.lower():
                 raise Tidy3dImportError(
                     "Module 'gdstk' not found. It is required to export shapes to gdstk cells."
                 )

tidy3d/components/medium.py

@@ -16,7 +16,6 @@
 import pydantic.v1 as pd
 import xarray as xr
 from numpy.typing import NDArray
-from scipy import signal
 
 from tidy3d.components.material.tcad.heat import ThermalSpecType
 from tidy3d.constants import (
@@ -3496,6 +3495,7 @@ def _real_partial_fraction_decomposition(
             ``tuple`` is an array of coefficients representing any direct polynomial term.
 
         """
+        from scipy import signal
 
         if a.ndim != 1 or np.any(np.iscomplex(a)):
             raise ValidationError(

tidy3d/components/mode/derivatives.py

@@ -3,13 +3,14 @@
 from __future__ import annotations
 
 import numpy as np
-import scipy.sparse as sp
 
 from tidy3d.constants import EPSILON_0, ETA_0
 
 
 def make_dxf(dls, shape, pmc):
     """Forward derivative in x."""
+    import scipy.sparse as sp
+
     Nx, Ny = shape
     if Nx == 1:
         return sp.csr_matrix((Ny, Ny))
@@ -23,6 +24,8 @@ def make_dxf(dls, shape, pmc):
 
 def make_dxb(dls, shape, pmc):
     """Backward derivative in x."""
+    import scipy.sparse as sp
+
     Nx, Ny = shape
     if Nx == 1:
         return sp.csr_matrix((Ny, Ny))
@@ -38,6 +41,8 @@ def make_dxb(dls, shape, pmc):
 
 def make_dyf(dls, shape, pmc):
     """Forward derivative in y."""
+    import scipy.sparse as sp
+
     Nx, Ny = shape
     if Ny == 1:
         return sp.csr_matrix((Nx, Nx))
@@ -51,6 +56,8 @@ def make_dyf(dls, shape, pmc):
 
 def make_dyb(dls, shape, pmc):
     """Backward derivative in y."""
+    import scipy.sparse as sp
+
     Nx, Ny = shape
     if Ny == 1:
         return sp.csr_matrix((Nx, Nx))
@@ -82,6 +89,7 @@ def create_s_matrices(omega, shape, npml, dls, eps_tensor, mu_tensor, dmin_pml=(
     """Makes the 'S-matrices'. When dotted with derivative matrices, they add
     PML. If dmin_pml is set to False, PML will not be applied on the "bottom"
     side of the domain."""
+    import scipy.sparse as sp
 
     # strip out some information needed
     Nx, Ny = shape

tidy3d/components/mode/solver.py

@@ -2,12 +2,9 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 import numpy as np
-import scipy.linalg as linalg
-import scipy.sparse as sp
-import scipy.sparse.linalg as spl
 
 from tidy3d.components.base import Tidy3dBaseModel
 from tidy3d.components.types import EpsSpecType, ModeSolverType, Numpy
@@ -30,6 +27,10 @@
 # Good conductor permittivity cut-off value. Let it be as large as possible so long as not causing overflow in
 # double precision. This value is very heuristic.
 GOOD_CONDUCTOR_CUT_OFF = 1e70
+
+if TYPE_CHECKING:
+    from scipy import sparse as sp
+
 # Consider a material to be good conductor if |ep| (or |mu|) > GOOD_CONDUCTOR_THRESHOLD * |pec_val|
 GOOD_CONDUCTOR_THRESHOLD = 0.9
 
@@ -453,6 +454,7 @@ def trim_small_values(cls, mat: sp.csr_matrix, tol: float) -> sp.csr_matrix:
         max_element = np.amax(np.abs(mat))
         mat.data *= np.logical_or(np.abs(mat.data) / max_element > tol, np.abs(mat.data) > tol)
         mat.eliminate_zeros()
+        return mat
 
     @classmethod
     def solver_diagonal(
@@ -468,6 +470,8 @@ def solver_diagonal(
         basis_E,
     ):
         """EM eigenmode solver assuming ``eps`` and ``mu`` are diagonal everywhere."""
+        import scipy.sparse as sp
+        import scipy.sparse.linalg as spl
 
         # code associated with these options is included below in case it's useful in the future
         enable_preconditioner = False
@@ -685,6 +689,7 @@ def solver_tensorial(
         cls, eps, mu, der_mats, num_modes, neff_guess, vec_init, mat_precision, direction
     ):
         """EM eigenmode solver assuming ``eps`` or ``mu`` have off-diagonal elements."""
+        import scipy.sparse as sp
 
         mode_solver_type = "tensorial"
         N = eps.shape[-1]
@@ -830,6 +835,7 @@ def solver_eigs(
             Number of eigenmodes to compute.
         guess_value : float, optional
         """
+        import scipy.sparse.linalg as spl
 
         values, vectors = spl.eigs(
             mat, k=num_modes, sigma=guess_value, tol=TOL_EIGS, v0=vec_init, M=M
@@ -868,6 +874,7 @@ def solver_eigs_relative(
             Number of eigenmodes to compute.
         guess_value : float, optional
         """
+        import scipy.linalg as linalg
 
         basis, _ = np.linalg.qr(basis_vecs)
         mat_basis = np.conj(basis.T) @ mat @ basis
@@ -878,22 +885,25 @@ def solver_eigs_relative(
 
     @classmethod
     def isinstance_complex(cls, vec_or_mat, tol=TOL_COMPLEX):
-        """Check if a numpy array or scipy csr_matrix has complex component by looking at
+        """Check if a numpy array or scipy.sparse.csr_matrix has complex component by looking at
         norm(x.imag)/norm(x)>TOL_COMPLEX
 
         Parameters
         ----------
         vec_or_mat : Union[np.ndarray, sp.csr_matrix]
         """
+        import scipy.sparse.linalg as spl
+        from scipy.sparse import csr_matrix
 
         if isinstance(vec_or_mat, np.ndarray):
             return np.linalg.norm(vec_or_mat.imag) / (np.linalg.norm(vec_or_mat) + fp_eps) > tol
-        if isinstance(vec_or_mat, sp.csr_matrix):
+        if isinstance(vec_or_mat, csr_matrix):
             mat_norm = spl.norm(vec_or_mat)
             mat_imag_norm = spl.norm(vec_or_mat.imag)
             return mat_imag_norm / (mat_norm + fp_eps) > tol
-
-        raise RuntimeError("Variable type should be either numpy array or scipy csr_matrix.")
+        raise RuntimeError(
+            f"Variable type should be either numpy array or scipy.sparse.csr_matrix, got {type(vec_or_mat)}."
+        )
 
     @classmethod
     def type_conversion(cls, vec_or_mat, new_dtype):

tidy3d/material_library/parametric_materials.py

@@ -13,13 +13,6 @@
 from tidy3d.constants import ELECTRON_VOLT, EPSILON_0, HBAR, K_B, KELVIN, Q_e
 from tidy3d.log import log
 
-try:
-    from scipy import integrate
-
-    INTEGRATE_AVAILABLE = True
-except ImportError:
-    INTEGRATE_AVAILABLE = False
-
 # default values of the physical parameters for graphene
 # scattering rate in eV
 GRAPHENE_DEF_GAMMA = 0.00041
@@ -233,6 +226,12 @@ def interband_conductivity(self, freqs: list[float]) -> list[complex]:
         List[complex]
             The list of corresponding interband conductivities, in S.
         """
+        try:
+            from scipy import integrate
+
+            INTEGRATE_AVAILABLE = True
+        except ImportError:
+            INTEGRATE_AVAILABLE = False
 
         def fermi(E: float) -> float:
             """Fermi distribution."""

tidy3d/plugins/autograd/primitives/interpolate.py

@@ -5,7 +5,6 @@
 import numpy as np
 from autograd.extend import defvjp, primitive
 from numpy.typing import NDArray
-from scipy.linalg import solve_banded
 
 from tidy3d.log import log
 
@@ -369,6 +368,8 @@ def _solve_tridiagonal(lower: NDArray, diag: NDArray, upper: NDArray, rhs: NDArr
     np.ndarray
         Solution vector
     """
+    from scipy.linalg import solve_banded
+
     n = diag.size
     ab = np.zeros((3, n))
     ab[0, 1:] = upper[:-1]

tidy3d/plugins/design/method.py

@@ -3,18 +3,18 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
 
 import numpy as np
 import pydantic.v1 as pd
-import scipy.stats.qmc as qmc
 
 from tidy3d.components.base import Tidy3dBaseModel
 from tidy3d.constants import inf
 
 from .parameter import ParameterAny, ParameterFloat, ParameterInt, ParameterType
 
-DEFAULT_MONTE_CARLO_SAMPLER_TYPE = qmc.LatinHypercube
+if TYPE_CHECKING:
+    from scipy.stats import qmc as qmc_type
 
 
 class Method(Tidy3dBaseModel, ABC):
@@ -786,14 +786,14 @@ class AbstractMethodRandom(MethodSample, ABC):
     )
 
     @abstractmethod
-    def _get_sampler(self, parameters: tuple[ParameterType, ...]) -> qmc.QMCEngine:
+    def _get_sampler(self, parameters: tuple[ParameterType, ...]) -> qmc_type.QMCEngine:
         """Sampler for this ``Method`` class. If ``None``, sets a default."""
 
     def _get_run_count(self, parameters: Optional[list] = None) -> int:
         """Return the maximum number of runs for the method based on current method arguments."""
         return self.num_points
 
-    def sample(self, parameters: tuple[ParameterType, ...], **kwargs) -> dict[str, Any]:
+    def sample(self, parameters: tuple[ParameterType, ...], **kwargs) -> list[dict[str, Any]]:
         """Defines how the design parameters are sampled on grid."""
 
         sampler = self._get_sampler(parameters)
@@ -823,11 +823,12 @@ class MethodMonteCarlo(AbstractMethodRandom):
     >>> method = tdd.MethodMonteCarlo(num_points=20)
     """
 
-    def _get_sampler(self, parameters: tuple[ParameterType, ...]) -> qmc.QMCEngine:
+    def _get_sampler(self, parameters: tuple[ParameterType, ...]) -> qmc_type.QMCEngine:
         """Sampler for this ``Method`` class."""
+        from scipy.stats import qmc
 
         d = len(parameters)
-        return DEFAULT_MONTE_CARLO_SAMPLER_TYPE(d=d, seed=self.seed)
+        return qmc.LatinHypercube(d=d, seed=self.seed)
 
 
 MethodType = Union[

tidy3d/plugins/dispersion/fit.py

@@ -8,7 +8,6 @@
 
 import numpy as np
 import requests
-import scipy.optimize as opt
 from pydantic.v1 import Field, validator
 from rich.progress import Progress
 
@@ -361,6 +360,7 @@ def _fit_single(
         Tuple[:class:`.PoleResidue`, float]
             Results of single fit: (dispersive medium, RMS error).
         """
+        import scipy.optimize as opt
 
         # NOTE: Not used
         def constraint(coeffs, _grad=None):

tidy3d/plugins/microwave/lobe_measurer.py

@@ -8,7 +8,6 @@
 import numpy as np
 import pydantic.v1 as pd
 from pandas import DataFrame
-from scipy.signal import find_peaks, peak_widths
 
 from tidy3d.components.base import Tidy3dBaseModel, cached_property, skip_if_fields_missing
 from tidy3d.components.types import ArrayFloat1D, ArrayLike, Ax
@@ -132,6 +131,8 @@ def lobe_measures(self) -> DataFrame:
         DataFrame
             A DataFrame containing all lobe measures, where rows indicate the lobe index.
         """
+        from scipy.signal import find_peaks
+
         if self.apply_cyclic_extension:
             angle, signal = self.cyclic_extension(self.angle, self.radiation_pattern)
         else:
@@ -231,6 +232,8 @@ def _calc_peak_widths(
         self, angle: ArrayLike, signal: ArrayLike, peaks: ArrayLike
     ) -> tuple[ArrayLike, ArrayLike, ArrayLike, ArrayLike]:
         """Get the peak widths in terms of the angular coordinates."""
+        from scipy.signal import peak_widths
+
         rel_height = 1.0 - self.width_measure
         last_element = len(signal) - 1
         left_ips = np.zeros_like(peaks)