Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jun 10, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions tidy3d/components/dispersion_fitter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion tidy3d/components/geometry/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
)
Expand Down
2 changes: 1 addition & 1 deletion tidy3d/components/medium.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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(
Expand Down
10 changes: 9 additions & 1 deletion tidy3d/components/mode/derivatives.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand All @@ -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))
Expand All @@ -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))
Expand All @@ -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))
Expand Down Expand Up @@ -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
Expand Down
26 changes: 18 additions & 8 deletions tidy3d/components/mode/solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand Down
13 changes: 6 additions & 7 deletions tidy3d/material_library/parametric_materials.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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."""
Expand Down
3 changes: 2 additions & 1 deletion tidy3d/plugins/autograd/primitives/interpolate.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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]
Expand Down
15 changes: 8 additions & 7 deletions tidy3d/plugins/design/method.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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[
Expand Down
2 changes: 1 addition & 1 deletion tidy3d/plugins/dispersion/fit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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):
Expand Down
5 changes: 4 additions & 1 deletion tidy3d/plugins/microwave/lobe_measurer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
3 changes: 2 additions & 1 deletion tidy3d/plugins/resonance/resonance.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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]
Expand Down