128 lines
4.5 KiB
Python
128 lines
4.5 KiB
Python
from typing import cast
|
|
|
|
import numpy
|
|
|
|
from ..fdfd import solvers
|
|
from ._solver_cases import solver_plumbing_case
|
|
from .utils import assert_close
|
|
|
|
|
|
def test_scipy_qmr_wraps_user_callback_without_recursion(monkeypatch) -> None:
|
|
seen: list[tuple[float, ...]] = []
|
|
|
|
def fake_qmr(a, b: numpy.ndarray, **kwargs):
|
|
kwargs['callback'](numpy.array([1.0, 2.0]))
|
|
return numpy.array([3.0, 4.0]), 0
|
|
|
|
monkeypatch.setattr(solvers.scipy.sparse.linalg, 'qmr', fake_qmr)
|
|
result = solvers._scipy_qmr(
|
|
solver_plumbing_case().a0,
|
|
numpy.array([1.0, 0.0]),
|
|
callback=lambda xk: seen.append(tuple(xk)),
|
|
)
|
|
|
|
assert_close(result, [3.0, 4.0])
|
|
assert seen == [(1.0, 2.0)]
|
|
|
|
|
|
def test_scipy_qmr_installs_logging_callback_when_missing(monkeypatch) -> None:
|
|
callback_seen: list[numpy.ndarray] = []
|
|
|
|
def fake_qmr(a, b: numpy.ndarray, **kwargs):
|
|
callback = kwargs['callback']
|
|
callback(numpy.array([5.0, 6.0]))
|
|
callback_seen.append(b.copy())
|
|
return numpy.array([7.0, 8.0]), 0
|
|
|
|
monkeypatch.setattr(solvers.scipy.sparse.linalg, 'qmr', fake_qmr)
|
|
result = solvers._scipy_qmr(solver_plumbing_case().a0, numpy.array([1.0, 0.0]))
|
|
|
|
assert_close(result, [7.0, 8.0])
|
|
assert len(callback_seen) == 1
|
|
|
|
|
|
def test_generic_forward_preconditions_system_and_guess(monkeypatch) -> None:
|
|
case = solver_plumbing_case()
|
|
captured: dict[str, numpy.ndarray | float | object] = {}
|
|
|
|
monkeypatch.setattr(solvers.operators, 'e_full', lambda *args, **kwargs: case.a0)
|
|
monkeypatch.setattr(solvers.operators, 'e_full_preconditioners', lambda dxes: (case.pl, case.pr))
|
|
|
|
def fake_solver(a, b: numpy.ndarray, **kwargs):
|
|
captured['a'] = a
|
|
captured['b'] = b
|
|
captured['x0'] = kwargs['x0']
|
|
captured['atol'] = kwargs['atol']
|
|
return case.solver_result
|
|
|
|
result = solvers.generic(
|
|
omega=case.omega,
|
|
dxes=case.dxes,
|
|
J=case.j,
|
|
epsilon=case.epsilon,
|
|
matrix_solver=fake_solver,
|
|
matrix_solver_opts={'atol': 1e-12},
|
|
E_guess=case.guess,
|
|
)
|
|
|
|
assert_close(cast(object, captured['a']).toarray(), (case.pl @ case.a0 @ case.pr).toarray()) # type: ignore[attr-defined]
|
|
assert_close(cast(numpy.ndarray, captured['b']), case.pl @ (-1j * case.omega * case.j))
|
|
assert_close(cast(numpy.ndarray, captured['x0']), case.pl @ case.guess)
|
|
assert captured['atol'] == 1e-12
|
|
assert_close(result, case.pr @ case.solver_result)
|
|
|
|
|
|
def test_generic_adjoint_preconditions_system_and_guess(monkeypatch) -> None:
|
|
case = solver_plumbing_case()
|
|
captured: dict[str, numpy.ndarray | float | object] = {}
|
|
|
|
monkeypatch.setattr(solvers.operators, 'e_full', lambda *args, **kwargs: case.a0)
|
|
monkeypatch.setattr(solvers.operators, 'e_full_preconditioners', lambda dxes: (case.pl, case.pr))
|
|
|
|
def fake_solver(a, b: numpy.ndarray, **kwargs):
|
|
captured['a'] = a
|
|
captured['b'] = b
|
|
captured['x0'] = kwargs['x0']
|
|
captured['rtol'] = kwargs['rtol']
|
|
return case.solver_result
|
|
|
|
result = solvers.generic(
|
|
omega=case.omega,
|
|
dxes=case.dxes,
|
|
J=case.j,
|
|
epsilon=case.epsilon,
|
|
matrix_solver=fake_solver,
|
|
matrix_solver_opts={'rtol': 1e-9},
|
|
E_guess=case.guess,
|
|
adjoint=True,
|
|
)
|
|
|
|
expected_matrix = (case.pl @ case.a0 @ case.pr).T.conjugate()
|
|
assert_close(cast(object, captured['a']).toarray(), expected_matrix.toarray()) # type: ignore[attr-defined]
|
|
assert_close(cast(numpy.ndarray, captured['b']), case.pr.T.conjugate() @ (-1j * case.omega * case.j))
|
|
assert_close(cast(numpy.ndarray, captured['x0']), case.pr.T.conjugate() @ case.guess)
|
|
assert captured['rtol'] == 1e-9
|
|
assert_close(result, case.pl.T.conjugate() @ case.solver_result)
|
|
|
|
|
|
def test_generic_without_guess_does_not_inject_x0(monkeypatch) -> None:
|
|
case = solver_plumbing_case()
|
|
captured: dict[str, object] = {}
|
|
|
|
monkeypatch.setattr(solvers.operators, 'e_full', lambda *args, **kwargs: case.a0)
|
|
monkeypatch.setattr(solvers.operators, 'e_full_preconditioners', lambda dxes: (case.pl, case.pr))
|
|
|
|
def fake_solver(a, b: numpy.ndarray, **kwargs):
|
|
captured['kwargs'] = kwargs
|
|
return numpy.array([1.0, -1.0])
|
|
|
|
result = solvers.generic(
|
|
omega=1.0,
|
|
dxes=case.dxes,
|
|
J=numpy.array([2.0, 3.0]),
|
|
epsilon=case.epsilon,
|
|
matrix_solver=fake_solver,
|
|
)
|
|
|
|
assert 'x0' not in cast(dict[str, object], captured['kwargs'])
|
|
assert_close(result, case.pr @ numpy.array([1.0, -1.0]))
|