From 41383de91017899aed1739b4dbd731e24919e7c6 Mon Sep 17 00:00:00 2001 From: Frederik Schubert Date: Thu, 5 Jun 2025 18:51:23 +0200 Subject: [PATCH] fix(tidy3d): lazy load scipy to reduce import time --- tidy3d/components/dispersion_fitter.py | 11 +++++--- tidy3d/components/geometry/base.py | 2 +- tidy3d/components/medium.py | 2 +- tidy3d/components/mode/derivatives.py | 10 ++++++- tidy3d/components/mode/solver.py | 26 +++++++++++++------ .../material_library/parametric_materials.py | 13 +++++----- .../autograd/primitives/interpolate.py | 3 ++- tidy3d/plugins/design/method.py | 15 ++++++----- tidy3d/plugins/dispersion/fit.py | 2 +- tidy3d/plugins/microwave/lobe_measurer.py | 5 +++- tidy3d/plugins/resonance/resonance.py | 3 ++- 11 files changed, 59 insertions(+), 33 deletions(-) diff --git a/tidy3d/components/dispersion_fitter.py b/tidy3d/components/dispersion_fitter.py index 04dd1e160..13258bd76 100644 --- a/tidy3d/components/dispersion_fitter.py +++ b/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 diff --git a/tidy3d/components/geometry/base.py b/tidy3d/components/geometry/base.py index 4e6e7d2b2..4d724d558 100644 --- a/tidy3d/components/geometry/base.py +++ b/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." ) diff --git a/tidy3d/components/medium.py b/tidy3d/components/medium.py index ade214f7c..d65aed8f3 100644 --- a/tidy3d/components/medium.py +++ b/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( diff --git a/tidy3d/components/mode/derivatives.py b/tidy3d/components/mode/derivatives.py index 2ff9cef1c..19f0f4b04 100644 --- a/tidy3d/components/mode/derivatives.py +++ b/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 diff --git a/tidy3d/components/mode/solver.py b/tidy3d/components/mode/solver.py index 37aefe57f..5755647a3 100644 --- a/tidy3d/components/mode/solver.py +++ b/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): diff --git a/tidy3d/material_library/parametric_materials.py b/tidy3d/material_library/parametric_materials.py index 3a1df7fcc..3f055091a 100644 --- a/tidy3d/material_library/parametric_materials.py +++ b/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.""" diff --git a/tidy3d/plugins/autograd/primitives/interpolate.py b/tidy3d/plugins/autograd/primitives/interpolate.py index 2e80e36b1..ac5b6c2d4 100644 --- a/tidy3d/plugins/autograd/primitives/interpolate.py +++ b/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] diff --git a/tidy3d/plugins/design/method.py b/tidy3d/plugins/design/method.py index ada58533d..22b0745bc 100644 --- a/tidy3d/plugins/design/method.py +++ b/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[ diff --git a/tidy3d/plugins/dispersion/fit.py b/tidy3d/plugins/dispersion/fit.py index 67bb6e584..b9977e08b 100644 --- a/tidy3d/plugins/dispersion/fit.py +++ b/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): diff --git a/tidy3d/plugins/microwave/lobe_measurer.py b/tidy3d/plugins/microwave/lobe_measurer.py index b55bf3ac3..c843314a8 100644 --- a/tidy3d/plugins/microwave/lobe_measurer.py +++ b/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) diff --git a/tidy3d/plugins/resonance/resonance.py b/tidy3d/plugins/resonance/resonance.py index 25ba8d5de..0f49c0f0c 100644 --- a/tidy3d/plugins/resonance/resonance.py +++ b/tidy3d/plugins/resonance/resonance.py @@ -6,7 +6,6 @@ from typing import Union import numpy as np -import scipy.linalg import xarray as xr from pydantic.v1 import Field, NonNegativeFloat, PositiveInt, validator @@ -361,6 +360,8 @@ def _solve_gen_eig_prob( I. Theory and application to a quantum-dynamics model," J. Chem. Phys. 102, 8001 (1995). """ + import scipy.linalg + eigvals_b, eigvecs_b = scipy.linalg.eig(b_matrix) large_inds = np.abs(eigvals_b) > rcond * np.amax(np.abs(eigvals_b)) eigvals_b = eigvals_b[large_inds]