42359882336

Committed 06 Nov 2026 11:57PM UTC coverage: 82.498% (+0.1%) from 82.355%

Build # 42359882336

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #1220 from wdhawkins/csda_wgdsa

Add support for WGDSA to CSDA problems.

Coverage Stats

129 of 131 new or added lines in 4 files covered. (98.47%)

155 existing lines in 7 files now uncovered.

32247 of 39088 relevant lines covered (82.5%)

82093451.16 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

93.65

/python/lib/solver.cc

// SPDX-FileCopyrightText: 2025 The OpenSn Authors <https://open-sn.github.io/opensn/>
// SPDX-License-Identifier: MIT

#include "python/lib/py_wrappers.h"
#include <pybind11/functional.h>
#include "framework/logging/log.h"
#include <pybind11/stl.h>
#include "framework/runtime.h"
#include "framework/field_functions/field_function_grid_based.h"
#include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/acceleration/discrete_ordinates_keigen_acceleration.h"
#include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/acceleration/cmfd_acceleration.h"
#include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/acceleration/scdsa_acceleration.h"
#include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/acceleration/smm_acceleration.h"
#include "modules/linear_boltzmann_solvers/discrete_ordinates_curvilinear_problem/discrete_ordinates_curvilinear_problem.h"
#include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/discrete_ordinates_problem.h"
#include "modules/linear_boltzmann_solvers/uncollided_problem/uncollided_problem.h"
#include "modules/linear_boltzmann_solvers/uncollided_problem/uncollided_solver.h"
#include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/io/discrete_ordinates_problem_io.h"
#include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/solvers/transient_solver.h"
#include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/solvers/steady_state_solver.h"
#include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/solvers/nl_keigen_solver.h"
#include "modules/linear_boltzmann_solvers/discrete_ordinates_problem/solvers/pi_keigen_solver.h"
#include "modules/linear_boltzmann_solvers/lbs_problem/io/lbs_problem_io.h"
#include "modules/linear_boltzmann_solvers/lbs_problem/lbs_problem.h"
#include "modules/linear_boltzmann_solvers/lbs_problem/compute/lbs_compute.h"
#include "modules/solver.h"
#include <pybind11/numpy.h>
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <map>
#include <memory>
#include <string>
#include <vector>

namespace opensn
{

namespace
{

/// Converts a balance table to the dictionary returned by ComputeBalanceTable.
py::dict
BalanceTableToDict(const BalanceTable& table)
{
  py::dict values;
  values["absorption_rate"] = table.absorption_rate;
  values["production_rate"] = table.production_rate;
  values["inflow_rate"] = table.inflow_rate;
  values["outflow_rate"] = table.outflow_rate;
  // The standard balance omits CSDA deposition, so it is misleading for CSDA runs.
  if (not table.csda_particle_balance.has_value())
    values["balance"] = table.balance;
  if (table.csda_particle_deposition_rate.has_value())
    values["csda_particle_deposition_rate"] = table.csda_particle_deposition_rate.value();
  if (table.csda_particle_balance.has_value())
    values["csda_particle_balance"] = table.csda_particle_balance.value();
  if (table.csda_energy_production_rate.has_value())
    values["csda_energy_production_rate"] = table.csda_energy_production_rate.value();
  if (table.csda_energy_inflow_rate.has_value())
    values["csda_energy_inflow_rate"] = table.csda_energy_inflow_rate.value();
  if (table.csda_energy_outflow_rate.has_value())
    values["csda_energy_outflow_rate"] = table.csda_energy_outflow_rate.value();
  if (table.csda_energy_balance.has_value())
    values["csda_energy_balance"] = table.csda_energy_balance.value();
  if (table.initial_inventory.has_value())
    values["initial_inventory"] = table.initial_inventory.value();
  if (table.final_inventory.has_value())
    values["final_inventory"] = table.final_inventory.value();
  if (table.predicted_inventory_change.has_value())
    values["predicted_inventory_change"] = table.predicted_inventory_change.value();
  if (table.actual_inventory_change.has_value())
    values["actual_inventory_change"] = table.actual_inventory_change.value();
  if (table.inventory_residual.has_value())
    values["inventory_residual"] = table.inventory_residual.value();
  return values;
}

} // namespace

// Wrap problem
void
WrapProblem(py::module& slv)
{
  // clang-format off
  // problem base
  auto problem = py::class_<Problem, std::shared_ptr<Problem>>(
    slv,
    "Problem",
    R"(
    Base class for all problems.

    Wrapper of :cpp:class:`opensn::Problem`.
    )"
  );
  // clang-format on
}

// Wrap solver
void
WrapSolver(py::module& slv)
{
  // clang-format off
  // solver base
  auto solver = py::class_<Solver, std::shared_ptr<Solver> >(
    slv,
    "Solver",
    R"(
    Base class for all solvers.

    Wrapper of :cpp:class:`opensn::Solver`.
    )"
  );
  solver.def(
    "Initialize",
    &Solver::Initialize,
    R"(
    Initialize the solver.

    Solvers check that the problem is in a supported state (for example, steady-state mode
    without external sources for the k-eigenvalue solvers) and raise ``ValueError``
    otherwise. Setup and post-setup validation must both succeed. Calling ``Initialize``
    again clears the previous initialized state first.

    For MPI transport solvers, this is a collective operation and must be called on all ranks.
    )"
  );
  solver.def(
    "Execute",
    &Solver::Execute,
    R"(
    Execute the solver.

    The problem's current state and the solver's requirements are checked first, so changes
    made to the problem after ``Initialize`` (for example a mode switch or ``SetXSMap``) are
    validated before they are used. ``Initialize`` must be called before ``Execute``.

    For MPI transport solvers, this is a collective operation and must be called on all ranks.
    )"
  );
  solver.def(
    "Advance",
    &Solver::Advance,
    R"(
    Advance the solver by one step.

    Changes made to the problem since the last step are checked as in ``Execute``.
    ``Initialize`` must be called first. Solver classes that do not implement stepping
    leave the state unchanged and log that ``Advance`` is unsupported.

    For MPI transport solvers, this is a collective operation and must be called on all ranks.
    )"
  );
  // clang-format on
}

// Wrap LBS solver
void
WrapLBS(py::module& slv)
{
  // clang-format off
  // LBS problem
  auto lbs_problem = py::class_<LBSProblem, std::shared_ptr<LBSProblem>, Problem>(
    slv,
    "LBSProblem",
    R"(
    Base class for all linear Boltzmann problems.

    Wrapper of :cpp:class:`opensn::LBSProblem`.
    )"
  );
  lbs_problem.def(
    "GetScalarFluxFieldFunction",
    [](LBSProblem& self, bool only_scalar_flux)
    {
      py::list field_function_list_per_group;
      for (unsigned int group = 0; group < self.GetNumGroups(); group++)
      {
        if (only_scalar_flux)
        {
          field_function_list_per_group.append(self.CreateScalarFluxFieldFunction(group, 0));
        }
        else
        {
          py::list field_function_list_per_moment;
          for (unsigned int moment = 0; moment < self.GetNumMoments(); moment++)
          {
            field_function_list_per_moment.append(self.CreateScalarFluxFieldFunction(group, moment));
          }
          field_function_list_per_group.append(field_function_list_per_moment);
        }
      }
      return field_function_list_per_group;
    },
    R"(
    Return scalar-flux or flux-moment field functions grouped by energy group.

    Parameters
    ----------
    only_scalar_flux : bool, default=True
        If True, returns only the zeroth moment (scalar flux) field function for each group.
        If False, returns all moment field functions for each group.

    Returns
    -------
    Union[List[pyopensn.fieldfunc.FieldFunctionGridBased], List[List[pyopensn.fieldfunc.FieldFunctionGridBased]]]
        If ``only_scalar_flux=True``:
        ``result[g]`` is the scalar-flux field function for group ``g``.

        If ``only_scalar_flux=False``:
        ``result[g][m]`` is the field function for group ``g`` and moment ``m``.

    Notes
    -----
    Field functions are created on demand from the current solver state.

    The returned field functions are snapshots of the solver state at creation time; they are not
    refreshed automatically if the solver state changes.

    They support ``CanUpdate()`` and ``Update()`` while their owning problem is still alive.
    Calling ``Update()`` explicitly refreshes the same field-function object from the solver's
    latest state.

    Calling ``GetScalarFluxFieldFunction(only_scalar_flux=False)`` creates all requested
    moments from the current ``phi`` iterate at the time of the call.

    In the nested form (``only_scalar_flux=False``), the moment index varies fastest
    within each group (inner index = moment, outer index = group).
    )",
    py::arg("only_scalar_flux") = true
  );
  lbs_problem.def(
    "CreateFieldFunction",
    static_cast<std::shared_ptr<FieldFunctionGridBased> (LBSProblem::*)(
      const std::string&, const std::string&, double)>(&LBSProblem::CreateFieldFunction),
    R"(
    Create a named scalar field function derived from a 1D XS or ``power``.

    Parameters
    ----------
    name : str
        Name to assign to the returned field function.
    xs_name : str
        Built-in 1D XS name, custom XS name, or the special value ``power``.
        CSDA-enabled discrete-ordinates problems also support
        ``csda_energy_deposition``, ``cepxs_energy_deposition``,
        ``csda_charge_deposition``,
        ``csda_charge_deposition_term``, and
        ``csda_charge_deposition_term_cellavg``. When CSDA is enabled,
        ``energy_deposition`` is an alias for the conservative
        ``csda_energy_deposition``. ``cepxs_energy_deposition`` preserves the
        imported CEPXS response plus the CSDA energy-space contribution.
    power_normalization_target : float, default=-1.0
        If positive, scale the derived field function so that the raw power field would
        integrate to this total power.

    Notes
    -----
    The returned field function is created on demand from the current scalar-flux iterate.
    For ordinary XS names this computes ``sum_g xs[g] * phi_g`` at each node.

    The returned field function is a snapshot of the solver state at creation time; it is not
    refreshed automatically if the solver state changes. It supports ``CanUpdate()`` and
    ``Update()`` while its owning problem is still alive. Calling ``Update()`` explicitly
    recomputes the same field function from the solver's latest state.

    If ``xs_name == "power"``, the same power-generation formula used elsewhere by the solver
    is applied on demand.

    CSDA deposition fields are derived from the current scalar flux and terminal
    charged-particle current information. They should normally be created after the
    CSDA solve has converged.

    If ``power_normalization_target > 0``, the returned field function is scaled using the power
    implied by the current scalar flux. This scaling affects only the returned field function;
    it does not mutate the solver's internal ``phi`` vectors.

    The returned field function is a fresh object created for this call. It is not
    automatically updated by later solves or timesteps.
    )",
    py::arg("name"),
    py::arg("xs_name"),
    py::arg("power_normalization_target") = -1.0
  );
  lbs_problem.def(
    "ComputeFieldFunctionPowerScaleFactor",
    &LBSProblem::ComputeFieldFunctionPowerScaleFactor,
    R"(
    Compute the scale factor that would normalize the current scalar flux's total power to a
    target value.

    This is a collective operation and must be called on all MPI ranks.

    Parameters
    ----------
    power_normalization_target : float
        Target total power (must be positive). The returned factor is
        ``power_normalization_target / global_total_power``, where ``global_total_power`` is the
        same MPI-reduced, kappa-weighted total power ``CreateFieldFunction(..., "power", ...)``
        and the ``power`` field function use.

    Notes
    -----
    This does not mutate any solver state. It is the same scale factor
    ``CreateFieldFunction(..., "power", power_normalization_target)`` applies internally, exposed
    directly so callers can apply it to their own derived quantities (e.g. postprocessor output)
    without needing a field function.
    )",
    py::arg("power_normalization_target")
  );
  lbs_problem.def(
    "GetTime",
    &LBSProblem::GetTime,
    R"(
    Get the current simulation time in seconds.
    )"
  );
  lbs_problem.def(
    "GetTimeStep",
    &LBSProblem::GetTimeStep,
    R"(
    Get the current timestep size.
    )"
  );
  lbs_problem.def(
    "ComputeFissionRate",
    [](LBSProblem& self, const std::string& scalar_flux_iterate)
    {
      const std::vector<double>* phi_ptr = nullptr;
      if (scalar_flux_iterate == "old")
      {
        phi_ptr = &self.GetPhiOldLocal();
      }
      else if (scalar_flux_iterate == "new")
      {
        phi_ptr = &self.GetPhiNewLocal();
      }
      else
      {
        throw std::invalid_argument("Unknown scalar_flux_iterate value: \"" + scalar_flux_iterate + "\".");
      }
      return ComputeFissionRate(self, *phi_ptr);
    },
    R"(
    Computes the total fission rate.

    Parameters
    ----------
    scalar_flux_iterate : {'old', 'new'}
        Specifies which scalar flux vector to use in the calculation.
            - 'old': Use the previous scalar flux iterate.
            - 'new': Use the current scalar flux iterate.

    Returns
    -------
    float
        The total fission rate, integrated over the physical volume. In RZ and 1D spherical
        geometry the integral covers the full revolution (factors of 2*pi and 4*pi).

    Raises
    ------
    ValueError
        If `scalar_flux_iterate` is not 'old' or 'new'.
    )",
    py::arg("scalar_flux_iterate")
  );
  lbs_problem.def(
    "ComputeFissionProduction",
    [](LBSProblem& self, const std::string& scalar_flux_iterate)
    {
      const std::vector<double>* phi_ptr = nullptr;
      if (scalar_flux_iterate == "old")
      {
        phi_ptr = &self.GetPhiOldLocal();
      }
      else if (scalar_flux_iterate == "new")
      {
        phi_ptr = &self.GetPhiNewLocal();
      }
      else
      {
        throw std::invalid_argument("Unknown scalar_flux_iterate value: \"" + scalar_flux_iterate + "\".");
      }
      return ComputeFissionProduction(self, *phi_ptr);
    },
    R"(
    Computes the total fission production (nu*fission).

    Parameters
    ----------
    scalar_flux_iterate : {'old', 'new'}
        Specifies which scalar flux vector to use in the calculation.
            - 'old': Use the previous scalar flux iterate.
            - 'new': Use the current scalar flux iterate.

    Returns
    -------
    float
        The total fission production, integrated over the physical volume. In RZ and 1D spherical
        geometry the integral covers the full revolution (factors of 2*pi and 4*pi).

    Raises
    ------
    ValueError
        If `scalar_flux_iterate` is not 'old' or 'new'.
    )",
    py::arg("scalar_flux_iterate")
  );
  lbs_problem.def(
    "GetPhiOldLocal",
    [](LBSProblem& self)
    {
      return convert_vector(self.GetPhiOldLocal());
    },
    R"(
    Get the previous scalar flux iterate (local vector).

    Returns
    -------
    memoryview
        Memory view of the local old scalar flux vector.
    )"
  );
  lbs_problem.def(
    "GetPhiNewLocal",
    [](LBSProblem& self)
    {
      return convert_vector(self.GetPhiNewLocal());
    },
    R"(
    Get the current scalar flux iterate (local vector).

    Returns
    -------
    memoryview
        Memory view of the local new scalar flux vector.
    )"
  );
  lbs_problem.def(
    "WriteFluxMoments",
    [](LBSProblem& self, const std::string& file_base)
    {
      LBSSolverIO::WriteFluxMoments(self, file_base);
    },
    R"(
    Write flux moments to file.

    Parameters
    ----------
    file_base: str
        File basename.
    )",
    py::arg("file_base")
  );
  lbs_problem.def(
    "CreateAndWriteSourceMoments",
    [](LBSProblem& self, const std::string& file_base)
    {
      std::vector<double> source_moments = self.MakeSourceMomentsFromPhi();
      LBSSolverIO::WriteFluxMoments(self, file_base, source_moments);
    },
    R"(
    Write source moments from latest flux iterate to file.

    Parameters
    ----------
    file_base: str
        File basename.
    )",
    py::arg("file_base")
  );
  lbs_problem.def(
    "ReadFluxMomentsAndMakeSourceMoments",
    [](LBSProblem& self, const std::string& file_base, bool single_file_flag)
    {
      std::vector<double> file_phi;
      LBSSolverIO::ReadFluxMoments(self, file_base, single_file_flag, file_phi);
      log.Log() << "Making source moments from flux file.";
      self.SetExtSrcMomentsFrom(self.MakeSourceMomentsFromPhi(file_phi));
    },
    R"(
    Read flux moments and compute corresponding source moments.

    Parameters
    ----------
    file_base: str
        File basename.
    single_file_flag: bool
        True to read the single file ``<file_base>.h5`` instead of the per-rank file
        ``<file_base><rank>.h5``. Serial runs only: a multi-rank run raises ``ValueError``
        before any file is opened.
    )",
    py::arg("file_base"),
    py::arg("single_file_flag")
  );
  lbs_problem.def(
    "ReadSourceMoments",
    [](LBSProblem& self, const std::string& file_base, bool single_file_flag)
    {
      auto ext_src_moments = self.GetExtSrcMomentsLocal();
      LBSSolverIO::ReadFluxMoments(self, file_base, single_file_flag, ext_src_moments);
      self.SetExtSrcMomentsFrom(ext_src_moments);
    },
    R"(
    Read source moments from file.

    Parameters
    ----------
    file_base: str
        File basename.
    single_file_flag: bool
        True to read the single file ``<file_base>.h5`` instead of the per-rank file
        ``<file_base><rank>.h5``. Serial runs only: a multi-rank run raises ``ValueError``
        before any file is opened.
    )",
    py::arg("file_base"),
    py::arg("single_file_flag")
  );
  lbs_problem.def(
    "ReadFluxMoments",
    [](LBSProblem& self, const std::string& file_base, bool single_file_flag)
    {
      LBSSolverIO::ReadFluxMoments(self, file_base, single_file_flag);
    },
    R"(
    Read flux moment data.

    Parameters
    ----------
    file_base: str
        File basename.
    single_file_flag: bool
        True to read the single file ``<file_base>.h5`` instead of the per-rank file
        ``<file_base><rank>.h5``. Serial runs only: a multi-rank run raises ``ValueError``
        before any file is opened.
    )",
    py::arg("file_base"),
    py::arg("single_file_flag")
  );
  lbs_problem.def(
    "WriteAngularFluxes",
    [](DiscreteOrdinatesProblem& self, const std::string& file_base)
    {
      DiscreteOrdinatesProblemIO::WriteAngularFluxes(self, file_base);
    },
    R"(
    Write angular flux data to file.

    Parameters
    ----------
    file_base: str
        File basename.
    )",
    py::arg("file_base")
  );
  lbs_problem.def(
    "ReadAngularFluxes",
    [](DiscreteOrdinatesProblem& self, const std::string& file_base)
    {
      DiscreteOrdinatesProblemIO::ReadAngularFluxes(self, file_base);
    },
    R"(
    Read angular fluxes from file.

    Parameters
    ----------
    file_base: str
        File basename.
    )",
    py::arg("file_base")
  );
  lbs_problem.def(
    "WriteSurfaceAngularFluxes",
    [](DiscreteOrdinatesProblem& self,
       const std::string& file_base,
       const std::vector<std::string>& boundary_surfaces,
       const std::map<std::string, std::pair<std::string, double>>& interior_surfaces)
    {
      DiscreteOrdinatesProblemIO::WriteSurfaceAngularFluxes(
        self, file_base, boundary_surfaces, interior_surfaces);
    },
    R"(
    Write surface angular flux data to file.

    Parameters
    ----------
    file_base: str
        File basename.
    boundary_surfaces: list[str], default=[]
        Boundary names to export.
    interior_surfaces: dict[str, tuple[str, float]], default={}
        Interior surfaces in the form {'name': ('axis', value)}, where axis is 'x', 'y', or 'z'.
        Each interior surface is written under two tags: '<name>_u' for faces whose outward normal
        aligns with the positive axis and '<name>_d' for faces with the opposite orientation.
        An interior surface must match at least one cell face, must lie on cell faces inside the
        mesh, and may not share a plane with another interior surface. Its generated tags may not
        equal a requested boundary name. It may not coincide with an exterior boundary; export
        that boundary by name using boundary_surfaces instead.

    Notes
    -----
    Requires ``options.save_angular_flux=True``. Collective over all ranks; surface selections may
    differ by rank. Each rank writes ``<file_base><rank>.h5``. The file records whether the problem
    is in adjoint mode. After an adjoint steady-state solve, the stored angular flux is the adjoint
    flux for the listed direction, so forward and adjoint data with the same indices refer to the
    same direction.
    )",
    py::arg("file_base"),
    py::arg("boundary_surfaces") = std::vector<std::string>{},
    py::arg("interior_surfaces") = std::map<std::string, std::pair<std::string, double>>{}
  );
  lbs_problem.def(
    "ReadSurfaceAngularFluxes",
    [](DiscreteOrdinatesProblem& self, const std::string& file_base, py::list surfaces)
    {
      std::vector<std::string> surface_ids;
      for (py::handle surface : surfaces)
        surface_ids.push_back(surface.cast<std::string>());

      const auto surface_fluxes =
        DiscreteOrdinatesProblemIO::ReadSurfaceAngularFluxes(self, file_base, surface_ids);
      py::list result;
      for (const auto& surface_flux : surface_fluxes)
      {
        py::dict mapping;
        mapping["cell_ids"] = surface_flux.mapping.cell_ids;
        mapping["num_face_nodes"] = surface_flux.mapping.num_face_nodes;
        mapping["cell_map"] = surface_flux.mapping.cell_map;
        mapping["centroid_origin"] = surface_flux.mapping.centroid_origin;
        mapping["centroid_spacing"] = surface_flux.mapping.centroid_spacing;
        mapping["cell_stride"] = surface_flux.mapping.cell_stride;
        mapping["nodes_x"] = surface_flux.mapping.nodes_x;
        mapping["nodes_y"] = surface_flux.mapping.nodes_y;
        mapping["nodes_z"] = surface_flux.mapping.nodes_z;

        py::dict data;
        data["omega"] = surface_flux.data.omega;
        data["mu"] = surface_flux.data.mu;
        data["wt_d"] = surface_flux.data.wt_d;
        data["M_ij"] = surface_flux.data.mass_matrix;
        data["fe_shape"] = surface_flux.data.fe_shape;
        data["psi"] = surface_flux.data.psi;
        data["node_index"] = surface_flux.data.node_index;
        data["dir_index"] = surface_flux.data.dir_index;

        py::dict entry;
        entry["groupset_id"] = surface_flux.groupset_id;
        entry["adjoint"] = surface_flux.adjoint;
        entry["surface_name"] = surface_flux.surface_name;
        entry["mapping"] = std::move(mapping);
        entry["data"] = std::move(data);
        result.append(std::move(entry));
      }
      return result;
    },
    R"(
    Read surface angular fluxes from file.

    Parameters
    ----------
    file_base: str
        File basename.
    surfaces: list[str]
        Stored surface tags to read. Use the boundary name for an exterior boundary. For an
        interior surface named '<name>', use '<name>_u' for faces whose outward normal aligns with
        the positive axis and '<name>_d' for faces with the opposite orientation.

    Returns
    -------
    List[dict]
        Surface mapping and angular-flux data for each groupset and requested surface. Each entry
        has keys 'groupset_id', 'adjoint', 'surface_name', 'mapping', and 'data'. 'adjoint' is
        True when the file was written from an adjoint problem. The 'cell_map' in 'mapping' maps
        face centroids, with 'centroid_origin' subtracted before each coordinate is divided by
        'centroid_spacing' (1e-9 times the global mesh extent) and rounded to an integer, to the
        surface face index. The same face has the same key from either side of an interior surface
        and in every file for the same mesh.

    Notes
    -----
    Each rank reads ``<file_base><rank>.h5``; the files must have been written with the same mesh
    partitioning.
    )",
    py::arg("file_base"),
    py::arg("surfaces")
  );
  lbs_problem.def(
    "SetPointSources",
    [](LBSProblem& self, py::kwargs& params)
    {
      // Collect the arguments first so that keyword order does not matter. The problem replaces
      // or extends its sources in one validated step.
      bool clear = false;
      std::vector<std::shared_ptr<PointSource>> sources;
      for (auto [key, value] : params)
      {
        auto c_key = key.cast<std::string>();
        if (c_key == "clear_point_sources")
          clear = value.cast<bool>();
        else if (c_key == "point_sources")
        {
          for (auto source : value.cast<py::list>())
            sources.push_back(source.cast<std::shared_ptr<PointSource>>());
        }
        else
          throw std::invalid_argument("SetPointSources: unknown argument '" + c_key +
                                      "'; expected 'point_sources' or 'clear_point_sources'.");
      }
      self.SetPointSources(sources, clear);
    },
    R"(
    Set or clear point sources.

    Parameters
    ----------
    clear_point_sources: bool, default=False
        If true, all current the point sources of the problem are deleted.
    point_sources: List[pyopensn.source.PointSource]
        List of new point sources to be added to the problem. When both arguments are given,
        the current sources are cleared before the new ones are added, regardless of keyword
        order. The new list is validated as a whole: a rejected call leaves the problem's
        sources unchanged.

    Notes
    -----
    This is collective. A source caches mesh data when added. Prefer a separate source object for
    each problem. Share one only when both problems use the same mesh and spatial layout.
    )"
  );
  lbs_problem.def(
    "SetVolumetricSources",
    [](LBSProblem& self, py::kwargs& params)
    {
      // Collect the arguments first so that keyword order does not matter. The problem replaces
      // or extends its sources in one validated step.
      bool clear = false;
      std::vector<std::shared_ptr<VolumetricSource>> sources;
      for (auto [key, value] : params)
      {
        auto c_key = key.cast<std::string>();
        if (c_key == "clear_volumetric_sources")
          clear = value.cast<bool>();
        else if (c_key == "volumetric_sources")
        {
          for (auto source : value.cast<py::list>())
            sources.push_back(source.cast<std::shared_ptr<VolumetricSource>>());
        }
        else
          throw std::invalid_argument("SetVolumetricSources: unknown argument '" + c_key +
                                      "'; expected 'volumetric_sources' or "
                                      "'clear_volumetric_sources'.");
      }
      self.SetVolumetricSources(sources, clear);
    },
    R"(
    Set or clear volumetric sources.

    Parameters
    ----------
    clear_volumetric_sources: bool, default=False
        If true, all current the volumetric sources of the problem are deleted.
    volumetric_sources: List[pyopensn.source.VolumetricSource]
        List of new volumetric sources to be added to the problem. When both arguments are given,
        the current sources are cleared before the new ones are added, regardless of keyword
        order. The new list is validated as a whole: a rejected call leaves the problem's
        sources unchanged.

    Notes
    -----
    This is collective. A source caches mesh data when added. Prefer a separate source object for
    each problem. Share one only when both problems use the same mesh and spatial layout.
    )"
  );
  lbs_problem.def(
    "SetXSMap",
    [](LBSProblem& self, py::kwargs& params)
    {
      BlockID2XSMap xs_map;
      for (auto [key, value] : params)
      {
        auto c_key = key.cast<std::string>();
        if (c_key == "xs_map")
        {
          auto xs_entries = value.cast<py::list>();
          for (auto entry : xs_entries)
          {
            InputParameters xs_entry_pars = LBSProblem::GetXSMapEntryBlock();
            xs_entry_pars.AssignParameters(pyobj_to_param_block("", entry.cast<py::dict>()));
            const auto& block_ids =
              xs_entry_pars.GetParam("block_ids").GetVectorValue<unsigned int>();
            auto xs = xs_entry_pars.GetSharedPtrParam<MultiGroupXS>("xs");
            for (const auto& block_id : block_ids)
              xs_map[block_id] = xs;
          }
        }
        else
          throw std::invalid_argument("SetXSMap: unknown argument '" + c_key +
                                      "'; expected 'xs_map'.");
      }
      self.SetBlockID2XSMap(xs_map);
    },
    R"(
    Replace the block-id to cross-section map.

    Parameters
    ----------
    xs_map: List[Dict]
        A list of block-id to cross-section mapping dictionaries. Each dictionary supports:
          - block_ids: List[int] (required)
              Mesh block ids to associate with the cross section.
          - xs: pyopensn.xs.MultiGroupXS (required)
              Cross section object.

    Raises
    ------
    ValueError
        If the problem was constructed with an ``uncollided_flux`` file, or if
        ``options.csda_enabled=True`` and the new cross sections split a
        charged-particle group block across groupsets, contain more than two
        charged-particle group ranges, have negative or non-finite stopping powers or
        non-positive energy-group widths, or have stopping power or energy
        bounds that do not match the problem's group count. CSDA also rejects
        conflicting supplied energy structures, unknown or conflicting charged-particle
        species, or a map with no energy structure.
        Individual materials may omit bounds and use the shared structure.

    Notes
    -----
    This is collective.

    The problem is refreshed immediately after replacing the map. Material metadata,
    precursor storage, GPU carriers, and derived solver state owned by the concrete
    problem are rebuilt for the new cross sections.

    If ``options.use_precursors=True``, this flag remains active across XS-map
    changes even when the current map has no precursor-bearing material. In that
    case delayed-neutron source terms are simply inactive until a later map provides
    precursor data again.

    Existing precursor concentrations are remapped by local cell and precursor-family
    index. When the new material for a cell has fewer precursor families, extra old
    families are dropped. When it has more families, newly introduced families are
    initialized to zero. If a cell changes through a material with zero precursors,
    the precursor history for that cell is discarded and any later reintroduced
    precursor families start from zero.

    If any fissionable material in the new map contains delayed-neutron precursor
    data and ``options.use_precursors=True``, all fissionable materials in the map
    must contain precursor data. Non-fissionable materials may have zero precursors.

    The new map is validated before it is committed. A validation error leaves the map unchanged.
    A later rebuild failure may leave partial changes. In time-dependent mode, every cross section
    must provide ``VELOCITY`` or ``INV_VELOCITY`` data.

    Forward/adjoint mode toggles via :meth:`LBSProblem.SetAdjoint` do not change this map.
    The ``MultiGroupXS`` objects themselves are mutable and shared by pointer, so a change to
    one, such as ``Scale``, is seen by every problem that uses it. The transport mode belongs
    to each problem: problems that share a ``MultiGroupXS`` handle can run in different modes.
    )"
  );
  lbs_problem.def(
    "ZeroPhi",
    [](LBSProblem& self)
    {
      self.ZeroPhi();
    },
    R"(
    Zero the scalar-flux vectors (``phi_old`` and ``phi_new``) in-place.
    )"
  );
  lbs_problem.def(
    "SetAdjoint",
    [](LBSProblem& self, bool adjoint)
    {
      self.SetAdjoint(adjoint);
    },
    py::arg("adjoint") = true,
    R"(
    Set forward/adjoint transport mode.

    Parameters
    ----------
    adjoint: bool, default=True
        ``True`` enables adjoint mode and ``False`` enables forward mode. Adjoint mode is
        supported only for Cartesian geometries with standard quadrature operators
        (``operator_method='standard'``), in steady-state mode, and without an
        uncollided flux. Uncollided flux generation (:class:`UncollidedProblem`) is
        forward only. An unsupported request raises ``ValueError`` before any state
        changes.

    Raises
    ------
    ValueError
        If ``adjoint=True`` and the problem is time-dependent, has
        ``options.csda_enabled=True``, was constructed with an
        ``uncollided_flux`` file, or is an :class:`UncollidedProblem`. The problem
        is left unchanged.

    Notes
    -----
    This is one of two supported mode-setting paths in Python:
      1. ``options={'adjoint': ...}`` in the problem constructor.
      2. ``SetAdjoint(...)`` (this method).

    If this changes mode, OpenSn performs a full mode-transition reset:
      - Materials are reinitialized in the selected mode.
      - Point and volumetric sources are cleared.
      - Boundary conditions are cleared.
      - Scalar and angular flux vectors are zeroed.

    If this is called with the same mode as the current setting, no reset is performed.

    The block-id to cross-section map is preserved across the transition.
    The mode change applies to this problem only. Other problems that share the mapped
    ``MultiGroupXS`` objects keep their own mode.

    After a mode change, reapply the desired driving terms before solving, typically:
      - :meth:`LBSProblem.SetPointSources`
      - :meth:`LBSProblem.SetVolumetricSources`
      - :meth:`DiscreteOrdinatesProblem.SetBoundaryOptions`

    This routine is intentionally destructive with respect to source/boundary/flux state
    to avoid hidden coupling between forward and adjoint setups.

    This is a collective operation and must be called on all ranks.
    )"
  );
  lbs_problem.def(
    "SetForward",
    &LBSProblem::SetForward,
    R"(
    Set forward transport mode.

    Equivalent to ``SetAdjoint(False)``.
    )"
  );
  lbs_problem.def(
    "IsAdjoint",
    &LBSProblem::IsAdjoint,
    R"(
    Return ``True`` if the problem is in adjoint mode, otherwise ``False``.
    )"
  );
  lbs_problem.def(
    "IsTimeDependent",
    &LBSProblem::IsTimeDependent,
    R"(
    Return ``True`` if the problem is in time-dependent mode, otherwise ``False``.
    )"
  );

  auto unc_problem =
    py::class_<UncollidedProblem, std::shared_ptr<UncollidedProblem>, LBSProblem>(
      slv,
      "UncollidedProblem",
      R"(
    Define an uncollided transport problem for a first-collision calculation.

    The problem stores the mesh, materials, point sources, boundaries, and
    near-source regions used by :class:`UncollidedSolver`. Its near-source
    ray-tracing and conservation treatment follow C. Woodsford, J. C. Ragusa,
    and J. E. Morel, "Sweep-based uncollided-flux treatment on unstructured
    grids," Progress in Nuclear Energy, vol. 200, 106524, 2026,
    https://doi.org/10.1016/j.pnucene.2026.106524.

    Wrapper of :cpp:class:`opensn::UncollidedProblem`.
    )");
  unc_problem.def(
    py::init(
      [](py::kwargs& params)
      {
        return UncollidedProblem::Create(kwargs_to_param_block(params));
      }),
    R"(
    Construct an uncollided transport problem.

    Parameters
    ----------
    mesh : Mesh
        Cartesian two- or three-dimensional spatial mesh.
    num_groups : int
        Number of energy groups.
    groupsets : List[Dict], default=[]
        Groupset definitions required by the common LBS problem interface. The
        uncollided calculation itself does not use an angular quadrature.
    xs_map : List[Dict], default=[]
        Block-to-cross-section mappings. Total cross sections are used for
        attenuation and are recorded in the output file for compatibility
        checking by the collided problem.
    boundary_conditions : List[Dict], default=[]
        Vacuum or reflecting boundary conditions, using the same ``name`` and
        ``type`` entries as ``DiscreteOrdinatesProblem``. Reflecting boundaries
        are represented with directly projected image sources and must be
        planar, mutually orthogonal symmetry planes without an opposing
        reflecting plane.
    point_sources : List[pyopensn.source.PointSource], default=[]
        Explicit point sources. At least one point source is required. For now,
        each source must lie strictly inside a single cell; sources exactly on
        faces, edges, or vertices are rejected.
    near_source : List[pyopensn.logvol.LogicalVolume]
        Near-source ray-traced region for each point source, in the same order
        as ``point_sources``. This list is required whenever point sources are
        provided. Its length must equal the number of point sources, and every
        point source must lie in its corresponding region.
    scattering_order : int, default=0
        Maximum spherical-harmonic order written to the HDF5 file. It must be
        at least the scattering order of the collided problem that consumes
        the file.
    Notes
    -----
    Only explicit point sources are supported by the uncollided generator.
    Each point source must lie strictly inside a single cell. Uncollided
    generation must run on a Cartesian two- or three-dimensional mesh with
    exactly one MPI rank.

    Point-source placement strongly affects spatial accuracy. Construct or
    locally refine the mesh so that each source lies well inside its containing
    cell and near the cell centroid whenever the physical model permits. Exact
    balance after the near-source correction does not guarantee that the local
    flux shape is resolved; monitor the reported correction and verify mesh
    convergence.

    A volumetric source may be approximated externally by multiple weighted
    point sources, as in the Kobayashi benchmark example. Choose the source
    quadrature and mesh together so that those points lie well inside their
    cells, and verify convergence with respect to both discretizations.

    Uncollided generation is forward only: ``options={'adjoint': True}`` and
    ``SetAdjoint(True)`` are rejected. The other ``options`` and ``use_gpus`` are
    inherited from the base class interface but have no effect on uncollided
    generation and should be omitted.

    The uncollided file generated by :class:`UncollidedSolver` is
    partition-independent and may be consumed by a serial or parallel collided
    calculation on the same mesh. The collided problem must use the same
    reflecting boundaries recorded in the file.
    Each reflected image source is ray traced to every finite-element volume
    quadrature point and projected directly into the spatial discretization.
    )");

  // discrete ordinate solver
  auto do_problem = py::class_<DiscreteOrdinatesProblem, std::shared_ptr<DiscreteOrdinatesProblem>,
                               LBSProblem>(
    slv,
    "DiscreteOrdinatesProblem",
    R"(
    Base class for discrete ordinates problems in Cartesian geometry.

    This class implements the algorithms necessary to solve a problem using the discrete ordinates method.
    When paired with a solver base, the result is a solver instance configured for a specific problem type
    (steady-state, transient, adjoint, k-eigenvalue, etc.).

    Wrapper of :cpp:class:`opensn::DiscreteOrdinatesProblem`.
    )"
  );
  do_problem.def(
    py::init(
      [](py::kwargs& params)
      {
        return DiscreteOrdinatesProblem::Create(kwargs_to_param_block(params));
      }
    ),
    R"(
    Construct a discrete ordinates problem with Cartesian geometry.

    Parameters
    ----------
    mesh : Mesh
        The spatial mesh.
    num_groups : int
        The total number of energy groups.
    groupsets : List[Dict]
        A list of input parameter blocks, each block provides the iterative properties for a
        groupset. The inclusive ranges must partition all groups in order, starting at group 0
        and ending at ``num_groups - 1``, without gaps or overlaps. Each dictionary supports:

          - groups_from_to: List[int] (required)
              Two-entry list with the first and last group id for the groupset, e.g. ``[0, 3]``.
          - angular_quadrature: pyopensn.aquad.AngularQuadrature, optional
              Handle to an angular quadrature.
          - angle_aggregation_type: {'polar', 'single', 'azimuthal', 'quadrant', 'octant'},
            default='single'
              Angle aggregation method to use during sweeping. ``'quadrant'`` (2D) and
              ``'octant'`` (3D) put all directions in the same quadrant or octant in one angle
              set. They require a Cartesian orthogonal mesh and a quadrature in which every
              direction has nonzero ``Omega_x`` and ``Omega_y`` (and, for ``'octant'``,
              ``Omega_z``), so each direction lies strictly inside one quadrant or octant.
          - inner_linear_method: {'classic_richardson', 'petsc_richardson',
            'petsc_gmres', 'petsc_fgmres', 'petsc_bicgstab'}, default='petsc_richardson'

              Iterative method used for inner linear solves. ``'petsc_fgmres'`` is flexible
              GMRES with right preconditioning; it remains exact when the WGDSA/TGDSA
              preconditioner solves its diffusion system only approximately.
          - l_abs_tol: float, default=1.0e-6
              Inner linear solver absolute residual tolerance.
          - l_max_its: int, default=200
              Inner linear solver maximum iterations.
          - gmres_restart_interval: int, default=30
              GMRES restart interval, if GMRES is used.
          - allow_cycles: bool, default=True
              Whether cyclic dependencies are allowed in sweeps.
          - apply_wgdsa: bool, default=False
              Enable within-group DSA for this groupset. With ``csda_enabled`` and charged-particle
              groups in the groupset, each group's diffusion operator also removes the CSDA
              slowing-down rate ``S_g / dE_g``, and the groups are solved one at a time in energy
              order, coupled by in-groupset scattering and slowing-down. See
              :ref:`theory_csda_wgdsa`.
          - wgdsa_l_abs_tol: float, default=1.0e-4
              Residual tolerance of the WGDSA diffusion solve, relative to the norm of its
              right-hand side (the name is historical).
          - wgdsa_l_max_its: int, default=30
              WGDSA maximum iterations.
          - wgdsa_verbose: bool, default=False
              Verbose WGDSA output.
          - wgdsa_petsc_options: str, default=''
              PETSc options string for the WGDSA solver.
          - wgdsa_solver_policy: {'auto', 'direct', 'iterative', 'petsc_options'}, default='auto'
              WGDSA diffusion solver policy. ``auto`` uses a direct PETSc LU solve below
              ``wgdsa_direct_solve_threshold`` global unknowns and an iterative solve
              otherwise. ``petsc_options`` lets ``wgdsa_petsc_options`` control the
              PETSc KSP/PC setup.
          - wgdsa_direct_solve_threshold: int, default=20000
              Maximum global WGDSA diffusion unknown count for the automatic direct solve. The
              CSDA operator solves one group at a time, so the count is that of one group.
          - wgdsa_boundary_closure: {'marshak', 'closure_r', 'dirichlet'}, default='marshak'
              Boundary condition of the WGDSA diffusion correction on vacuum, isotropic and
              arbitrary boundaries. ``marshak`` imposes the natural Marshak condition, a face
              mass term with coefficient 1/2; ``closure_r`` blends the
              Marshak form written with the quadrature's partial current (face mass
              coefficient 2 beta_f, beta_f the quadrature's half-range moment, against the
              analytic 1/2) on thin boundary cells with the half-range transport form on
              thick ones, uses beta_f as the floor of the interior penalty, and is faster
              than ``marshak`` when boundary cells are many mean free paths thick;
              ``dirichlet`` (OpenSn's original) imposes a zero correction weakly
              (interior-penalty term). Reflecting boundaries are zero-current.
          - apply_tgdsa: bool, default=False
              Enable two-grid DSA for this groupset.
          - tgdsa_l_abs_tol: float, default=1.0e-4
              Residual tolerance of the TGDSA diffusion solve, relative to the norm of its
              right-hand side (the name is historical).
          - tgdsa_l_max_its: int, default=30
              TGDSA maximum iterations.
          - tgdsa_verbose: bool, default=False
              Verbose TGDSA output.
          - tgdsa_petsc_options: str, default=''
              PETSc options string for the TGDSA solver.
          - tgdsa_solver_policy: {'auto', 'direct', 'iterative', 'petsc_options'}, default='auto'
              TGDSA diffusion solver policy. ``auto`` uses a direct PETSc LU solve below
              ``tgdsa_direct_solve_threshold`` global unknowns and an iterative solve
              otherwise. ``petsc_options`` lets ``tgdsa_petsc_options`` control the
              PETSc KSP/PC setup.
          - tgdsa_direct_solve_threshold: int, default=20000
              Maximum global TGDSA diffusion unknown count for the automatic direct solve.
          - tgdsa_boundary_closure: {'marshak', 'closure_r', 'dirichlet'}, default='marshak'
              Boundary condition of the TGDSA diffusion correction, with the same choices
              and boundaries as ``wgdsa_boundary_closure``.
          - wgdsa_h_perp_convention: {'face_side', 'cell'}, default='face_side'
              Penalty length of the WGDSA MIP operator. ``face_side`` uses the face-dependent
              orthogonal length (3V/A_f tet, V/A_f hex and quad, 2V/A_f triangle) capped by four
              times the smallest height of the face's PWLD side simplices, with an extra factor 2
              on polygons and polyhedra; the cap keeps the operator positive definite on strongly
              distorted cells and is inactive on regular ones. ``cell`` (OpenSn's original) uses
              the cell-wide lengths 3V/S (tet) and V/S (hex) with the extra factor 2, and on
              polygons the face-dependent lengths without the cap or the factor 2. On polyhedra
              its penalty is larger (six times on cubes), which slows DSA on cells a few mean free
              paths thick; on polygons it is half the ``face_side`` penalty and can make the
              operator indefinite on strongly distorted cells.
          - tgdsa_h_perp_convention: {'face_side', 'cell'}, default='face_side'
              Penalty length of the TGDSA MIP operator, as ``wgdsa_h_perp_convention``.
    xs_map : List[Dict]
        A list of mappings from block ids to cross-section definitions. Each dictionary supports:
          - block_ids: List[int] (required)
              Mesh block IDs to associate with the cross section.
          - xs: pyopensn.xs.MultiGroupXS (required)
              Cross-section object to assign to the specified blocks.
    boundary_conditions: List[Dict], default=[]
        A list containing tables for each boundary specification. Each dictionary supports:
          - name: str (required)
              Boundary name that identifies the specific boundary.
          - type: {'vacuum', 'isotropic', 'reflecting', 'arbitrary'} (required)
              Boundary type specification.
          - group_strength: List[float], optional
              Required when ``type='isotropic'``. Isotropic strength per group.
          - start_time: float, optional
              Active start time for isotropic boundaries only. Defaults to -infinity.
          - end_time: float, optional
              Active end time for isotropic boundaries only. Defaults to infinity.
          - function: AngularFluxFunction, optional
              Required when ``type='arbitrary'`` unless ``time_function`` is supplied. Callable
              that returns incoming angular flux from group and direction.
          - time_function: AngularFluxTimeFunction, optional
              Required when ``type='arbitrary'`` unless ``function`` is supplied. Callable that
              returns incoming angular flux from group, direction, and time.

        Isotropic boundaries may use ``start_time``/``end_time`` for simple on/off behavior.
        Arbitrary boundaries must specify exactly one of ``function`` or ``time_function``; use
        ``time_function`` for time-dependent arbitrary inflow and handle any active window inside
        that callback.
    point_sources: List[pyopensn.source.PointSource], default=[]
        A list of point sources.
    volumetric_sources: List[pyopensn.source.VolumetricSource], default=[]
        A list of volumetric sources.
    options : Dict, default={}
        A block of optional configuration parameters, including:
          - max_mpi_message_size: int, default=4194304
              Maximum size in bytes of one sweep MPI message. Larger angular-flux transfers
              are split into messages of at most this size.
          - restart_writes_enabled: bool, default=False
              Enable restart dump writes for solvers that support restart output.
          - write_delayed_psi_to_restart: bool, default=True
              Include delayed sweep angular-flux buffers in steady-state restart dumps. These
              buffers are optional when a steady-state restart is used only as a transient
              initial condition because the transient initialization can reconstruct angular
              state from the flux moments. Time-dependent restart dumps always include them,
              regardless of this option, because continuing a transient requires them whenever
              the problem has delayed sweep angular state, including partitioned parallel,
              reflected-boundary, and cyclic-sweep cases.
          - write_angular_flux_to_restart: bool, default=True
              Include stored angular fluxes in steady-state restart dumps when
              ``save_angular_flux=True``. They are optional when a steady-state restart is used
              only as a transient initial condition. Time-dependent restart dumps always
              include them, regardless of this option, because continuing a transient requires
              them.
          - read_restart_path: str, default=''
              File stem for reading a full restart. Scalar flux and precursor data can be
              remapped when the MPI rank count or mesh partition changes. Full transient
              continuation still requires the original partitioned angular state layout.
          - read_initial_condition_path: str, default=''
              File stem for reading restart data as an initial condition. Cannot be combined
              with ``read_restart_path``. A steady-state restart may be used by
              ``TransientSolver`` through this option. The first ``TransientSolver.Initialize()``
              loads it once; later initializations continue from the problem's current state.
          - write_restart_path: str, default=''
              File stem for restart dump writes. OpenSn appends the MPI rank and
              ``.restart.h5`` to this stem.
          - write_restart_time_interval: int, default=0
            (must be 0 or >=30)
          - use_precursors: bool, default=True
            Enable delayed neutrons. With ``False`` every solve (steady-state, k-eigenvalue and
            transient) uses prompt neutrons only, even if the cross sections have delayed-neutron
            data; a k-eigenvalue solve then gives the prompt eigenvalue. This is treated as user
            intent and remains active across later ``SetXSMap`` calls, even if the current XS map
            temporarily contains no precursor-bearing material. When XS maps are swapped, existing
            precursor concentrations are remapped by cell and precursor-family index; new families
            start at zero and removed families are discarded. If any fissionable material in the
            active map has precursor data and ``use_precursors=True``, all fissionable materials in
            that map must have precursor data. Non-fissionable materials may have zero precursors.
          - delayed_neutrons: bool, default=True
            Alias of ``use_precursors``. Give at most one of the two names.
          - use_source_moments: bool, default=False
            Use the problem's source moments as the fixed source instead of its
            volumetric and point sources. Provide them (e.g. with ``ReadSourceMoments``)
            after creating the problem and before solving; a solve without them raises
            ``ValueError``. They are an external source, so the k-eigenvalue solvers
            reject a problem with this option.
          - save_angular_flux: bool, default=False
            Store angular flux state (`psi`) for transient mode, angular-flux
            field functions, and angular-flux I/O.
          - deterministic_sweeps: bool, default=False
            Make every sweep bitwise reproducible from run to run: the scalar flux moments and
            face outflows are summed on the host from stored angular fluxes in a fixed angle order
            instead of in the order that MPI messages and GPU threads complete. Stores the angular
            flux in host memory without enabling ``save_angular_flux``. GPU sweeps also copy it to
            the host after every sweep, which makes them roughly 2-5 times slower. Meant for
            debugging and reproducible comparisons.
          - adjoint: bool, default=False
          - csda_enabled: bool, default=False
            Enable CSDA charged-particle transport for CEPXS data loaded with
            ``csda_format=True``. CSDA requires a Cartesian
            :class:`SteadyStateSourceSolver` solve, ``sweep_type='AAH'``, CPU
            sweeps, no ``apply_tgdsa`` in any groupset, and charged-particle
            group ranges that are not split across groupsets. ``apply_wgdsa``
            uses a CSDA diffusion operator in groupsets with charged-particle
            groups. All materials must supply identical per-group energy bounds
            or omit them; at least one material must supply the shared structure.
          - verbose_inner_iterations: bool, default=True
            Print inner iteration details, including WGS and AGS iterations.
          - verbose_outer_iterations: bool, default=True
            Print outer solver progress, including PI/NLKE iterations and transient steps.
          - max_ags_iterations: int, optional
            Maximum across-groupset iterations per transport solve. A value set here is used by
            every solver. If it is not set, each solver uses its own default: 100 for the
            steady-state and transient solvers, and 1 in each power iteration of
            :class:`pyopensn.solver.PowerIterationKEigenSolver`.
          - ags_tolerance: float, default=1.0e-6
          - ags_convergence_check: {'l2', 'pointwise'}, default='l2'
          - power_default_kappa: float, default=3.20435e-11
          - field_function_prefix_option: {'prefix', 'solver_name'}, default='prefix'
          - field_function_prefix: str, default=''
          - fused_aao_sweep: bool, default=False
            Use the fused all-at-once GPU scheduler for ``AAH`` sweeps. Each kernel launch
            sweeps the current level of up to 256 ready angle sets instead of one, which
            reduces launch count when individual levels are too small to fill the GPU. Up to
            ``fused_aao_lanes`` such runs are in flight at once, so angle sets that become
            ready later (for example when upstream angular fluxes arrive from another rank)
            start without waiting for the others. The solution matches the default scheduler
            to floating-point rounding. Requires ``sweep_type="AAH"`` and ``use_gpus=True``.
            Otherwise construction raises an error.
          - fused_aao_lanes: int, default=8
            With ``fused_aao_sweep=True``, the number of fused runs (each a batch of angle
            sets on its own GPU stream) that may be in flight at once on each rank. Must be at
            least 1. Ignored unless ``fused_aao_sweep=True``.

        These options are applied at problem creation.
    sweep_type : str, default="AAH"
        The sweep type to use. Must be one of `AAH` or `CBC`. Defaults to `AAH`.
        Both sweep types support time-dependent (transient) mode.
    time_dependent : bool, default=False
        If true, the problem starts in time-dependent mode. Otherwise it starts in
        steady-state mode. Requires ``options.save_angular_flux=True`` and
        ``VELOCITY`` or ``INV_VELOCITY`` data in every cross section.
        Both ``AAH`` and ``CBC`` sweep types support time-dependent mode.
    uncollided_flux : str, default=""
        HDF5 file generated by :class:`UncollidedSolver`. For steady-state
        forward fixed-source calculations, OpenSn uses its moments to
        construct the first-collision scattering and fission source, solves
        for the collided component, and adds the uncollided moments back to
        the converged state. The file must match the current group count, mesh
        cell IDs, cell-node layout and coordinates, total cross sections, and
        reflecting boundary set. Its maximum moment order must be at least
        this problem's scattering order. The same serial file may be read by
        any MPI partitioning of the matching mesh. The first-collision source
        replaces the fixed sources, so the problem must not also have volumetric
        or point sources or incoming-flux (isotropic or arbitrary) boundaries,
        and its cross sections cannot be replaced.
    use_gpus : bool, default=False
        A flag specifying whether GPU acceleration is used for the sweep. Both
        ```AAH``` and ```CBC``` sweep types support GPU acceleration.

    Notes
    -----
    Construction validates and builds the runtime problem. It is collective. Use equivalent input
    on every MPI rank.
    )"
  );
  do_problem.def(
    "SetTimeDependentMode",
    py::overload_cast<>(&DiscreteOrdinatesProblem::SetTimeDependentMode),
    R"(
    Set the problem to time-dependent mode.

    Notes
    -----
    Switch problem from steady-state to time-dependent mode. This updates problem
    internals (sweep chunk mode and source-function) while preserving user boundary
    conditions and fixed sources.

    Requires ``options.save_angular_flux=True`` at problem creation and ``VELOCITY`` or
    ``INV_VELOCITY`` data in every cross section. Not supported for adjoint, GPU,
    curvilinear, or uncollided-flux problems. An unsupported request raises
    ``ValueError`` before any state changes.

    This is a collective operation and must be called on all ranks.
    )"
  );
  do_problem.def(
    "SetSteadyStateMode",
    &DiscreteOrdinatesProblem::SetSteadyStateMode,
    R"(
    Set the problem to steady-state mode.

    Notes
    -----
    Switch problem from time-dependent to steady-state mode. This updates problem
    internals (sweep chunk mode and source-function) while preserving user boundary
    conditions and fixed sources.

    This is a collective operation and must be called on all ranks.
    )"
  );
  do_problem.def(
    "IsTimeDependent",
    &DiscreteOrdinatesProblem::IsTimeDependent,
    R"(
    Return ``True`` if the problem is currently in time-dependent mode.
    )"
  );
  do_problem.def(
    "SetBoundaryOptions",
    [](DiscreteOrdinatesProblem& self, py::kwargs& params)
    {
      bool clear_boundary_conditions = false;
      std::vector<InputParameters> boundary_params;
      for (auto [key, value] : params)
      {
        auto c_key = key.cast<std::string>();
        if (c_key == "clear_boundary_conditions")
          clear_boundary_conditions = value.cast<bool>();
        else if (c_key == "boundary_conditions")
        {
          auto boundaries = value.cast<py::list>();
          for (auto boundary : boundaries)
          {
            InputParameters input = DiscreteOrdinatesProblem::GetBoundaryOptionsBlock();
            input.AssignParameters(pyobj_to_param_block("", boundary.cast<py::dict>()));
            boundary_params.push_back(std::move(input));
          }
        }
        else
          throw std::invalid_argument("SetBoundaryOptions: unknown argument '" + c_key +
                                      "'; expected 'boundary_conditions' or "
                                      "'clear_boundary_conditions'.");
      }
      if (clear_boundary_conditions or not boundary_params.empty())
        self.SetBoundaryOptions(boundary_params, clear_boundary_conditions);
    },
    R"(
    Set or clear boundary conditions.

    Parameters
    ----------
    clear_boundary_conditions: bool, default=False
        If true, all current boundary conditions are deleted.
    boundary_conditions: List[Dict]
        A list of boundary condition dictionaries. Each dictionary supports:
          - name: str (required)
              Boundary name that identifies the specific boundary.
          - type: {'vacuum', 'isotropic', 'reflecting', 'arbitrary'} (required)
              Boundary type specification.
          - group_strength: List[float], optional
              Required when ``type='isotropic'``. Isotropic strength per group.
          - start_time: float, optional
              Active start time for isotropic boundaries only. Defaults to -infinity.
          - end_time: float, optional
              Active end time for isotropic boundaries only. Defaults to infinity.
          - function: AngularFluxFunction, optional
              Required when ``type='arbitrary'`` unless ``time_function`` is supplied. Callable
              that returns incoming angular flux from group and direction.
          - time_function: AngularFluxTimeFunction, optional
              Required when ``type='arbitrary'`` unless ``function`` is supplied. Callable that
              returns incoming angular flux from group, direction, and time.

        Isotropic boundaries may use ``start_time``/``end_time`` for simple on/off behavior.
        Arbitrary boundaries must specify exactly one of ``function`` or ``time_function``; use
        ``time_function`` for time-dependent arbitrary inflow and handle any active window inside
        that callback.

    Notes
    -----
    Mode transitions via :meth:`LBSProblem.SetAdjoint` clear all boundary conditions.
    Reapply boundaries with this method before solving in the new mode.

    This is a collective operation and must be called on all ranks.
    )"
  );
  do_problem.def(
    "ZeroPsi",
    [](DiscreteOrdinatesProblem& self)
    {
      self.ZeroPsi();
    },
    R"(
    Zero the angular-flux storage.
    )"
  );
  do_problem.def(
    "GetPsi",
    [](DiscreteOrdinatesProblem& self)
    {
      const auto& psi = self.GetPsiNewLocal();
      py::list psi_list;
      for (const auto& vec : psi)
      {
        auto array = py::array_t<double>(static_cast<py::ssize_t>(vec.size()));
        std::copy(vec.begin(), vec.end(), static_cast<double*>(array.mutable_data()));
        psi_list.append(array);
      }
      return psi_list;
    },
    R"(
    Return psi as a list of NumPy arrays (float64).

    The arrays are copies of the current angular-flux state. Mutating them does not
    mutate the problem.
    )"
  );
  do_problem.def(
    "GetAngularFieldFunctionList",
    [](DiscreteOrdinatesProblem& self, py::list groups, py::list angles)
    {
      std::vector<unsigned int> group_ids;
      std::vector<size_t> angle_ids;
      group_ids.reserve(groups.size());
      angle_ids.reserve(angles.size());

      for (py::handle g : groups)
        group_ids.push_back(g.cast<unsigned int>());
      for (py::handle a : angles)
        angle_ids.push_back(a.cast<size_t>());

      auto ff_list = self.CreateAngularFluxFieldFunctionList(group_ids, angle_ids);
      py::list out;
      for (const auto& ff : ff_list)
        out.append(ff);
      return out;
    },
    R"(
    Create field functions for selected angular flux components.

    Note: You must enable angular flux storage (``save_angular_flux=True``) in
    the problem options. For transient problems this is required. The angular
    flux is also stored when ``deterministic_sweeps=True``. If neither option
    is set, the returned field functions will remain zero.

    Example
    -------
    .. code-block:: python

       solver.Initialize()
       solver.Execute()
       ang_ff = phys.GetAngularFieldFunctionList(groups=[0], angles=[0])

    For transient/time-dependent runs, each field function is still a snapshot. Either call
    this after each timestep to create a fresh object or keep the returned object and call
    ``Update()`` after each timestep before exporting or interpolating it. Two common patterns
    are:

    1) Use ``TransientSolver.Execute()`` with a post-advance callback:

    .. code-block:: python

       solver = TransientSolver(problem=phys)
       solver.Initialize()
       ang_ff = phys.GetAngularFieldFunctionList(groups=[0], angles=[0])

       def post_advance():
           for ff in ang_ff:
               ff.Update()
           FieldFunctionGridBased.ExportMultipleToPVTU(ang_ff, "angular_flux_t")

       solver.SetPostAdvanceCallback(post_advance)
       solver.Execute()

    2) Use a custom Python loop with ``TransientSolver.Advance()``:

    .. code-block:: python

       solver = TransientSolver(problem=phys)
       solver.Initialize()
       ang_ff = phys.GetAngularFieldFunctionList(groups=[0], angles=[0])
       for _ in range(num_steps):
           solver.Advance()
           for ff in ang_ff:
               ff.Update()
           FieldFunctionGridBased.ExportMultipleToPVTU(ang_ff, "angular_flux_t")

    Parameters
    ----------
    groups : List[int]
        Global group indices to export.
    angles : List[int]
        Angle indices within the groupset quadrature for each group.

    Returns
    -------
    List[pyopensn.fieldfunc.FieldFunctionGridBased]
        Field functions for the requested ``(group, angle)`` pairs. Each returned field function
        is a snapshot, but supports ``CanUpdate()`` and ``Update()`` while its owning problem is
        alive.
    )",
    py::arg("groups"),
    py::arg("angles")
  );
  do_problem.def(
    "ComputeLeakage",
    [](DiscreteOrdinatesProblem& self, py::list bnd_names)
    {
      auto grid = self.GetGrid();
      std::map<std::uint64_t, std::string> allowed_bd_ids = grid->GetBoundaryIDMap();
      const auto coord_sys = grid->GetCoordinateSystem();
      const auto mesh_type = grid->GetType();
      const auto dim = grid->GetDimension();
      // Resolve the requested boundaries, preserving user order. An empty list selects all.
      std::vector<std::uint64_t> bndry_ids;
      if (not bnd_names.empty())
      {
        for (py::handle name : bnd_names)
        {
          const auto sname = name.cast<std::string>();
          const auto bid = self.FindBoundaryID(sname);
          if (not bid.has_value())
            throw std::invalid_argument("ComputeLeakage: Boundary name '" + sname +
                                        "' is not valid for this problem.");
          bndry_ids.push_back(*bid);
        }
      }
      else
      {
        bndry_ids = self.GetGrid()->GetUniqueBoundaryIDs();
      }
      // compute the leakage
      std::map<std::uint64_t, std::vector<double>> leakage = ComputeLeakage(self, bndry_ids);
      const bool rz_ortho = (coord_sys == CoordinateSystemType::CYLINDRICAL &&
                             mesh_type == MeshType::ORTHOGONAL && dim == 2);

      auto ToRZName = [](const std::string& name)
      {
        if (name == "xmin") return std::string("rmin");
        if (name == "xmax") return std::string("rmax");
        if (name == "ymin") return std::string("zmin");
        if (name == "ymax") return std::string("zmax");
        return name;
      };

      // convert result to native Python
      py::dict result;
      for (const auto& bndry_id : bndry_ids)
      {
        const auto it = leakage.find(bndry_id);
        if (it == leakage.end())
          continue;
        // construct numpy array and copy contents
        const auto& grp_wise_leakage = it->second;
        py::array_t<double> np_vector(static_cast<py::ssize_t>(grp_wise_leakage.size()));
        auto buffer = np_vector.request();
        auto *np_vector_data = static_cast<double*>(buffer.ptr);
        std::copy(grp_wise_leakage.begin(), grp_wise_leakage.end(), np_vector_data);
        std::string name = allowed_bd_ids.at(bndry_id);
        if (rz_ortho)
          name = ToRZName(name);
        result[py::str(name)] = std::move(np_vector);
      }

      return result;
    },
    R"(
    Compute leakage for the problem.

    Parameters
    ----------
    bnd_names : List[str]
        A list of boundary names for which leakage should be computed. An empty list selects all
        boundaries. RZ problems use rmin, rmax, zmin, and zmax.

    Returns
    -------
    Dict[str, numpy.ndarray]
        A dictionary mapping boundary names to group-wise leakage vectors.
        Each array contains the net leakage (per group): the outgoing partial current minus the
        incoming partial current, integrated over the corresponding boundary surface. On a
        vacuum boundary this equals the outgoing current; on a boundary with an incident flux it
        is reduced by the inflow and can be negative; on a reflecting boundary it vanishes at
        convergence. In RZ geometry the surface integral covers the full revolution.

    Raises
    ------
    ValueError
        If one or more boundary ids are not present on the current mesh.
    )",
    py::arg("bnd_names")
  );

  // discrete ordinates curvilinear problem
  auto do_curvilinear_problem = py::class_<DiscreteOrdinatesCurvilinearProblem,
                                           std::shared_ptr<DiscreteOrdinatesCurvilinearProblem>,
                                           DiscreteOrdinatesProblem>(
    slv,
    "DiscreteOrdinatesCurvilinearProblem",
    R"(
    Base class for discrete ordinates problems in curvilinear geometry.

    Wrapper of :cpp:class:`opensn::DiscreteOrdinatesCurvilinearProblem`.
    )"
  );
  do_curvilinear_problem.def(
    py::init(
      [](py::kwargs& params)
      {
        return DiscreteOrdinatesCurvilinearProblem::Create(kwargs_to_param_block(params));
      }
    ),
    R"(
    Construct a discrete ordinates problem for curvilinear geometry.

    Two geometries are supported: 2D cylindrical (``r-z``, a 2D mesh with
    ``coord_sys="cylindrical"`` and :class:`pyopensn.aquad.GLCProductQuadrature2DRZ`) and 1D
    spherical (a 1D radial mesh with ``coord_sys="spherical"`` and
    :class:`pyopensn.aquad.GLProductQuadrature1DSpherical`; boundaries ``zmin`` (inner) and
    ``zmax`` (outer), and a mesh starting at r = 0 is a solid sphere whose center needs no
    boundary condition). Groupsets must use ``angle_aggregation_type`` ``'azimuthal'`` or
    ``'single'``, ``sweep_type`` must be ``'AAH'``, and time-dependent problems and GPU
    acceleration are not supported. Integrated quantities are reported for the full revolution
    (``r-z``) or the full sphere.

    Warnings
    --------
       DiscreteOrdinatesCurvilinearProblem is **experimental** and should be used with caution!

    Parameters
    ----------
    mesh : Mesh
        The spatial mesh. Its geometry type (cylindrical or spherical) is taken directly from
        the mesh -- e.g. from the ``coord_sys`` argument used to build it with a mesh
        generator. There is no separate ``coord_system`` parameter on this class; supplying
        one is rejected as an unknown parameter.
    num_groups : int
        The total number of energy groups.
    groupsets : list of dict
        A list of input parameter blocks, each block provides the iterative properties for a
        groupset. Each dictionary supports:

          - groups_from_to: List[int] (required)
              Two-entry list with the first and last group id for the groupset, e.g. ``[0, 3]``.
          - angular_quadrature: pyopensn.aquad.AngularQuadrature, optional
              Handle to an angular quadrature.
          - angle_aggregation_type: {'single', 'azimuthal'}, default='single'
              Angle aggregation method to use during sweeping. Curvilinear problems do not
              support ``'polar'``.
          - inner_linear_method: {'classic_richardson', 'petsc_richardson',
            'petsc_gmres', 'petsc_fgmres', 'petsc_bicgstab'}, default='petsc_richardson'

              Iterative method used for inner linear solves. ``'petsc_fgmres'`` is flexible
              GMRES with right preconditioning; it remains exact when the WGDSA/TGDSA
              preconditioner solves its diffusion system only approximately.
          - l_abs_tol: float, default=1.0e-6
              Inner linear solver absolute residual tolerance.
          - l_max_its: int, default=200
              Inner linear solver maximum iterations.
          - gmres_restart_interval: int, default=30
              GMRES restart interval, if GMRES is used.
          - allow_cycles: bool, default=True
              Whether cyclic dependencies are allowed in sweeps.
          - apply_wgdsa: bool, default=False
              Enable within-group DSA for this groupset.
          - wgdsa_l_abs_tol: float, default=1.0e-4
              Residual tolerance of the WGDSA diffusion solve, relative to the norm of its
              right-hand side (the name is historical).
          - wgdsa_l_max_its: int, default=30
              WGDSA maximum iterations.
          - wgdsa_verbose: bool, default=False
              Verbose WGDSA output.
          - wgdsa_petsc_options: str, default=''
              PETSc options string for the WGDSA solver.
          - wgdsa_solver_policy: {'auto', 'direct', 'iterative', 'petsc_options'}, default='auto'
              WGDSA diffusion solver policy. ``auto`` uses a direct PETSc LU solve below
              ``wgdsa_direct_solve_threshold`` global unknowns and an iterative solve
              otherwise. ``petsc_options`` lets ``wgdsa_petsc_options`` control the
              PETSc KSP/PC setup.
          - wgdsa_direct_solve_threshold: int, default=20000
              Maximum global WGDSA diffusion unknown count for the automatic direct solve.
          - wgdsa_boundary_closure: {'marshak', 'closure_r', 'dirichlet'}, default='marshak'
              Boundary condition of the WGDSA diffusion correction on vacuum, isotropic and
              arbitrary boundaries. ``marshak`` imposes the natural Marshak condition, a face
              mass term with coefficient 1/2; ``closure_r`` blends the
              Marshak form written with the quadrature's partial current (face mass
              coefficient 2 beta_f, beta_f the quadrature's half-range moment, against the
              analytic 1/2) on thin boundary cells with the half-range transport form on
              thick ones, uses beta_f as the floor of the interior penalty, and is faster
              than ``marshak`` when boundary cells are many mean free paths thick;
              ``dirichlet`` (OpenSn's original) imposes a zero correction weakly
              (interior-penalty term). Reflecting boundaries are zero-current.
          - apply_tgdsa: bool, default=False
              Enable two-grid DSA for this groupset.
          - tgdsa_l_abs_tol: float, default=1.0e-4
              Residual tolerance of the TGDSA diffusion solve, relative to the norm of its
              right-hand side (the name is historical).
          - tgdsa_l_max_its: int, default=30
              TGDSA maximum iterations.
          - tgdsa_verbose: bool, default=False
              Verbose TGDSA output.
          - tgdsa_petsc_options: str, default=''
              PETSc options string for the TGDSA solver.
          - tgdsa_solver_policy: {'auto', 'direct', 'iterative', 'petsc_options'}, default='auto'
              TGDSA diffusion solver policy. ``auto`` uses a direct PETSc LU solve below
              ``tgdsa_direct_solve_threshold`` global unknowns and an iterative solve
              otherwise. ``petsc_options`` lets ``tgdsa_petsc_options`` control the
              PETSc KSP/PC setup.
          - tgdsa_direct_solve_threshold: int, default=20000
              Maximum global TGDSA diffusion unknown count for the automatic direct solve.
          - tgdsa_boundary_closure: {'marshak', 'closure_r', 'dirichlet'}, default='marshak'
              Boundary condition of the TGDSA diffusion correction, with the same choices
              and boundaries as ``wgdsa_boundary_closure``.
          - wgdsa_h_perp_convention: {'face_side', 'cell'}, default='face_side'
              Penalty length of the WGDSA MIP operator. ``face_side`` uses the face-dependent
              orthogonal length (3V/A_f tet, V/A_f hex and quad, 2V/A_f triangle) capped by four
              times the smallest height of the face's PWLD side simplices, with an extra factor 2
              on polygons and polyhedra; the cap keeps the operator positive definite on strongly
              distorted cells and is inactive on regular ones. ``cell`` (OpenSn's original) uses
              the cell-wide lengths 3V/S (tet) and V/S (hex) with the extra factor 2, and on
              polygons the face-dependent lengths without the cap or the factor 2. On polyhedra
              its penalty is larger (six times on cubes), which slows DSA on cells a few mean free
              paths thick; on polygons it is half the ``face_side`` penalty and can make the
              operator indefinite on strongly distorted cells.
          - tgdsa_h_perp_convention: {'face_side', 'cell'}, default='face_side'
              Penalty length of the TGDSA MIP operator, as ``wgdsa_h_perp_convention``.
    xs_map : list of dict
        A list of mappings from block ids to cross-section definitions. Each dictionary supports:
          - block_ids: List[int] (required)
              Mesh block IDs to associate with the cross section.
          - xs: pyopensn.xs.MultiGroupXS (required)
              Cross-section object to assign to the specified blocks.
    boundary_conditions: List[Dict], default=[]
        A list containing tables for each boundary specification. Each dictionary supports:
          - name: str (required)
              Boundary name that identifies the specific boundary.
          - type: {'vacuum', 'isotropic', 'reflecting', 'arbitrary'} (required)
              Boundary type specification.
          - group_strength: List[float], optional
              Required when ``type='isotropic'``. Isotropic strength per group.
          - start_time: float, optional
              Active start time for isotropic boundaries only. Defaults to -infinity.
          - end_time: float, optional
              Active end time for isotropic boundaries only. Defaults to infinity.
          - function: AngularFluxFunction, optional
              Required when ``type='arbitrary'`` unless ``time_function`` is supplied. Callable
              that returns incoming angular flux from group and direction.
          - time_function: AngularFluxTimeFunction, optional
              Required when ``type='arbitrary'`` unless ``function`` is supplied. Callable that
              returns incoming angular flux from group, direction, and time.

        Isotropic boundaries may use ``start_time``/``end_time`` for simple on/off behavior.
        Arbitrary boundaries must specify exactly one of ``function`` or ``time_function``; use
        ``time_function`` for time-dependent arbitrary inflow and handle any active window inside
        that callback.
    point_sources: List[pyopensn.source.PointSource], default=[]
        A list of point sources.
    volumetric_sources: List[pyopensn.source.VolumetricSource], default=[]
        A list of volumetric sources.
    options : dict, optional
        A block of optional configuration parameters applied at problem creation, including:
          - max_mpi_message_size: int, default=4194304
              Maximum size in bytes of one sweep MPI message. Larger angular-flux transfers
              are split into messages of at most this size.
          - restart_writes_enabled: bool, default=False
              Enable restart dump writes for solvers that support restart output.
          - write_delayed_psi_to_restart: bool, default=True
              Include delayed sweep angular-flux buffers in steady-state restart dumps. These
              buffers are optional when a steady-state restart is used only as a transient
              initial condition because the transient initialization can reconstruct angular
              state from the flux moments. Time-dependent restart dumps always include them,
              regardless of this option, because continuing a transient requires them whenever
              the problem has delayed sweep angular state, including partitioned parallel,
              reflected-boundary, and cyclic-sweep cases.
          - write_angular_flux_to_restart: bool, default=True
              Include stored angular fluxes in steady-state restart dumps when
              ``save_angular_flux=True``. They are optional when a steady-state restart is used
              only as a transient initial condition. Time-dependent restart dumps always
              include them, regardless of this option, because continuing a transient requires
              them.
          - read_restart_path: str, default=''
              File stem for reading a full restart. Scalar flux and precursor data can be
              remapped when the MPI rank count or mesh partition changes. Full transient
              continuation still requires the original partitioned angular state layout.
          - read_initial_condition_path: str, default=''
              File stem for reading restart data as an initial condition. Cannot be combined
              with ``read_restart_path``. A steady-state restart may be used by
              ``TransientSolver`` through this option. The first ``TransientSolver.Initialize()``
              loads it once; later initializations continue from the problem's current state.
          - write_restart_path: str, default=''
              File stem for restart dump writes. OpenSn appends the MPI rank and
              ``.restart.h5`` to this stem.
          - write_restart_time_interval: int, default=0
          - use_precursors: bool, default=True
            Enable delayed neutrons. With ``False`` every solve uses prompt neutrons only, even if
            the cross sections have delayed-neutron data. This remains active across later
            ``SetXSMap`` calls, even if the current map temporarily has no precursor-bearing
            material. When XS maps change, precursor concentrations are remapped by cell and
            precursor-family index; newly introduced families start at zero and removed families
            are discarded.
          - delayed_neutrons: bool, default=True
            Alias of ``use_precursors``. Give at most one of the two names.
          - use_source_moments: bool, default=False
            Use the problem's source moments as the fixed source instead of its
            volumetric and point sources. Provide them (e.g. with ``ReadSourceMoments``)
            after creating the problem and before solving; a solve without them raises
            ``ValueError``. They are an external source, so the k-eigenvalue solvers
            reject a problem with this option.
          - save_angular_flux: bool, default=False
            Store angular flux state (`psi`) for transient mode, angular-flux
            field functions, and angular-flux I/O.
          - deterministic_sweeps: bool, default=False
            Make every sweep bitwise reproducible from run to run: the scalar flux moments and
            face outflows are summed from stored angular fluxes in a fixed angle order instead of
            in the order that MPI messages arrive. Stores the angular flux in host memory without
            enabling ``save_angular_flux``. Meant for debugging and reproducible comparisons.
          - csda_enabled: bool, default=False
            Enable CSDA charged-particle transport for CEPXS data loaded with
            ``csda_format=True``. CSDA is not supported for this curvilinear
            problem class.
          - verbose_inner_iterations: bool, default=True
            Print inner iteration details, including WGS and AGS iterations.
          - verbose_outer_iterations: bool, default=True
            Print outer solver progress, including PI/NLKE iterations and transient steps.
          - max_ags_iterations: int, optional
            Maximum across-groupset iterations per transport solve. A value set here is used by
            every solver. If it is not set, each solver uses its own default: 100 for the
            steady-state and transient solvers, and 1 in each power iteration of
            :class:`pyopensn.solver.PowerIterationKEigenSolver`.
          - ags_tolerance: float, default=1.0e-6
          - ags_convergence_check: {'l2', 'pointwise'}, default='l2'
          - power_default_kappa: float, default=3.20435e-11
          - field_function_prefix_option: {'prefix', 'solver_name'}, default='prefix'
          - field_function_prefix: str, default=''
    sweep_type : str, optional
        The sweep type to use. Curvilinear problems support only `AAH`, the default.
        Time-dependent mode is not supported for curvilinear problems.
    )"
  );
}

// Wrap uncollided solver
void
WrapUncollidedSolver(py::module& slv)
{
  // clang-format off
  auto uncollided_solver =
    py::class_<UncollidedSolver, std::shared_ptr<UncollidedSolver>, Solver>(
      slv,
      "UncollidedSolver",
      R"(
    Generate and write uncollided flux moments for first-collision transport.

    Wrapper of :cpp:class:`opensn::UncollidedSolver`.
    )"
  );
  uncollided_solver.def(
    py::init(
      [](py::kwargs& params)
      {
        return UncollidedSolver::Create(kwargs_to_param_block(params));
      }
    ),
    R"(
    Construct an uncollided transport solver.

    Parameters
    ----------
    problem : pyopensn.solver.UncollidedProblem
        Existing uncollided problem instance.
    file_name : str
        Output HDF5 file containing uncollided flux moments, mesh and total
        cross-section metadata, and production/removal/outflow balance terms.
    progress_interval : int, default=5
        Percentage interval for source-point progress reports. Reports include
        completed source points, elapsed time, and estimated remaining time.
        Set to zero to disable progress reporting. Valid values are integers in
        ``[0, 100]``.

    Notes
    -----
    :meth:`Initialize` requires exactly one MPI rank. The generated HDF5 file
    may then be reused by a serial or parallel collided calculation on the
    same mesh. :meth:`Execute` uses the problem's point sources as they are when
    it runs, including sources replaced after :meth:`Initialize`.

    Internal threading in the current uncollided implementation is capped by
    the ``OPENSN_NUM_THREADS`` environment variable and defaults to ``1`` when
    the variable is unset or invalid.
    )"
  );
  // clang-format on
}

// Wrap steady-state solver
void
WrapSteadyState(py::module& slv)
{

  // clang-format off
  // steady state solver
  auto steady_state_solver = py::class_<SteadyStateSourceSolver, std::shared_ptr<SteadyStateSourceSolver>,
                                        Solver>(
    slv,
    "SteadyStateSourceSolver",
    R"(
    Steady state solver.

    Wrapper of :cpp:class:`opensn::SteadyStateSourceSolver`.
    )"
  );
  steady_state_solver.def(
    py::init(
      [](py::kwargs& params)
      {
        return SteadyStateSourceSolver::Create(kwargs_to_param_block(params));
      }
    ),
    R"(
    Construct a steady state solver.

    Parameters
    ----------
    problem : pyopensn.solver.DiscreteOrdinatesProblem
        Existing discrete ordinates problem instance.
    compute_balance : bool, optional, default=False
        If true, print the balance table after :meth:`Execute` completes.

    Notes
    -----
    If the problem was constructed with ``options={'read_restart_path': ...}``,
    restart data is read during :meth:`Initialize`. If it was constructed with
    ``options={'restart_writes_enabled': True, ...}``, a restart dump is written
    after :meth:`Execute` completes.

    Restart files are rank-layout specific. A restart written with ``N`` MPI
    ranks must be read with the same rank count and a compatible problem
    definition.
    )"
  );
  steady_state_solver.def(
    "GetNumSweeps",
    &SteadyStateSourceSolver::GetNumSweeps,
    R"(
    Return the cumulative number of transport sweeps performed by all groupsets.

    Returns
    -------
    int
        Total number of transport sweeps.
    )"
  );
  steady_state_solver.def(
    "IsConverged",
    &SteadyStateSourceSolver::IsConverged,
    R"(
    Return True if the transport solve of the last :meth:`Execute` converged.

    With more than one groupset this is the across-groupset iteration, otherwise the
    within-group solve of the single groupset. Returns False if the iteration limit was
    reached or the solve failed, if the solve applied no convergence test (a single sweep,
    ``petsc_richardson`` with ``l_max_its = 1``), and before the first :meth:`Execute`.
    )"
  );
  steady_state_solver.def(
    "ComputeBalanceTable",
    [](const SteadyStateSourceSolver& self)
    {
      return BalanceTableToDict(self.ComputeBalanceTable());
    },
    R"(
    Compute and return the global balance table using the solver's normalization.
    This is a collective operation and must be called on all ranks.

    Returns
    -------
    dict
        Dictionary with the following entries:

        - ``absorption_rate``:
          Global absorption rate, approximately ``integral sigma_a * phi dV`` summed over
          groups and the full domain.
        - ``production_rate``:
          Global volumetric production/source rate used by the solver,
          approximately ``integral Q dV`` summed over groups and the full domain.
        - ``inflow_rate``:
          Global incoming boundary contribution integrated over incoming
          angular flux on boundaries. On reflecting boundaries the outgoing partial
          current is counted as the incoming one, so reflecting faces add equal
          amounts to ``inflow_rate`` and ``outflow_rate`` and cancel in ``balance``.
        - ``outflow_rate``:
          Global outgoing boundary contribution accumulated from face outflow
          tallies, including reflecting boundaries (see ``inflow_rate``).
        - ``balance``:
          Rate balance,
          ``production_rate + inflow_rate - absorption_rate - outflow_rate``.
          Omitted when CSDA is enabled, because it does not include CSDA
          particle deposition; use ``csda_particle_balance`` instead.
        - ``csda_particle_deposition_rate``:
          Present only when CSDA is enabled. Rate at which particles slow down
          out of the lowest-energy group of each charged-particle block. Electron
          and positron contributions are both counted as positive.
        - ``csda_particle_balance``:
          Present only when CSDA is enabled. Signed relative particle residual:
          ``production_rate + inflow_rate - absorption_rate - outflow_rate -
          csda_particle_deposition_rate``, divided by
          ``production_rate + inflow_rate``. Zero gain gives zero for a zero
          residual and signed infinity otherwise.
        - ``csda_energy_production_rate``, ``csda_energy_inflow_rate``, and
          ``csda_energy_outflow_rate``:
          Present only when CSDA is enabled. Production, boundary inflow, and
          boundary outflow rates, each weighted by group-midpoint energy. For a
          converged solve, ``csda_energy_production_rate +
          csda_energy_inflow_rate - csda_energy_outflow_rate`` equals the
          integral of the ``csda_energy_deposition`` field function.
        - ``csda_energy_balance``:
          Present only when CSDA is enabled. Signed relative energy residual:
          energy production plus boundary inflow, minus boundary outflow,
          collision energy loss, and CSDA continuous energy loss, all weighted by
          group-midpoint energy, divided by
          ``csda_energy_production_rate + csda_energy_inflow_rate``. Zero gain
          gives zero for a zero residual and signed infinity otherwise.

    Notes
    -----
    This solver applies no extra normalization to the balance table.
    CSDA balance entries are convergence-sensitive and should be interpreted only
    after the linear solve has converged to the requested tolerance.
    )"
  );
  // clang-format on
}

// Wrap transient solver
void
WrapTransient(py::module& slv)
{
  // clang-format off
  auto transient_solver =
    py::class_<TransientSolver, std::shared_ptr<TransientSolver>, Solver>(
      slv,
      "TransientSolver",
      R"(
      Transient solver.

      Wrapper of :cpp:class:`opensn::TransientSolver`.
      )"
    );
  transient_solver.def(
    py::init(
      [](py::kwargs& params)
      {
        return TransientSolver::Create(kwargs_to_param_block(params));
      }
    ),
    R"(
    Construct a transient solver.

    Parameters
    ----------
    pyopensn.solver.DiscreteOrdinatesProblem : DiscreteOrdinatesProblem
        Existing discrete ordinates problem instance.
    dt : float, optional, default=2.0e-3
        Time step size used during the simulation.
    stop_time : float, optional, default=0.1
        Absolute end time of the transient, not a duration. A solver that starts after
        another one on the same problem continues from that problem's time.
    theta : float, optional, default=0.5
        Time differencing scheme parameter.
    initial_state : str, optional, default="existing"
        Initial state for the transient solve. Allowed values: existing, zero.
        "existing" continues from the problem's current solution and time. "zero"
        resets the scalar flux, precursors, and stored angular fluxes to zero and starts
        at time 0. Restart data and a pending initial condition
        (``read_initial_condition_path``) take precedence.
    verbose : bool, optional, default=True
        Enable verbose logging.
    compute_balance : bool, optional, default=False
        If true, print the timestep balance table (the one :meth:`ComputeBalanceTable`
        returns) after every timestep, whether taken by :meth:`Execute` or :meth:`Advance`.

    Notes
    -----
    The associated problem must have ``save_angular_flux=True`` enabled. This
    is required for transient problems.

    If the problem was constructed with ``options={'read_restart_path': ...}``,
    restart data is read during :meth:`Initialize`. If it was constructed with
    ``options={'read_initial_condition_path': ...}``, restart data is read as a
    transient initial condition during :meth:`Initialize`, then the problem is
    switched to time-dependent mode; a steady-state restart may be used in this
    mode. In this initial-condition path, stored angular fluxes and delayed
    sweep angular-flux buffers are optional in the steady-state restart; the
    transient solver reconstructs angular state from the loaded flux moments.
    If it was constructed with
    ``options={'restart_writes_enabled': True, ...}``, timed restart dumps may
    be written during :meth:`Execute` and a final restart dump is written when
    execution completes.

    A full transient continuation restart read with ``read_restart_path`` is
    different from a steady-state initial-condition restart read with
    ``read_initial_condition_path``. Full transient continuation restarts require
    angular flux state in the restart file; if the problem has delayed sweep
    angular state, including partitioned parallel, reflected-boundary, or
    cyclic-sweep cases, continuation also requires delayed sweep angular-flux
    buffers in the restart. Restart files are rank-layout specific: use the same MPI rank
    count and compatible problem definition when reading files written by a
    previous run.
    )"
  );
  transient_solver.def(
    "SetTimeStep",
    &TransientSolver::SetTimeStep,
    R"(
    Set the timestep size used by :meth:`Advance`.

    Parameters
    ----------
    dt : float
        New timestep size.
    )");
  transient_solver.def(
    "SetTheta",
    &TransientSolver::SetTheta,
    R"(
    Set the theta parameter used by :meth:`Advance`.

    Parameters
    ----------
    theta : float
        Theta value between 1.0e-16 and 1.
    )");
  transient_solver.def(
    "Advance",
    &TransientSolver::Advance,
    R"(
    Advance the solver by a single timestep.

    Notes
    -----
    You must call :meth:`Initialize` before calling :meth:`Advance` or
    :meth:`Execute`.
    )");
  transient_solver.def(
    "SetPreAdvanceCallback",
    static_cast<void (TransientSolver::*)(std::function<void()>)>(
      &TransientSolver::SetPreAdvanceCallback),
    R"(
    Register a callback that runs before each advance within :meth:`Execute`.

    :param callback: Function invoked before the solver advances a timestep. Pass None
        to clear. If the callback modifies the timestep, the new value is used for the
        upcoming step.
    :type callback: Optional[Callable[[], None]]
    )");
  transient_solver.def(
    "SetPreAdvanceCallback",
    py::overload_cast<std::nullptr_t>(&TransientSolver::SetPreAdvanceCallback),
    "Clear the PreAdvance callback by passing None.");
  transient_solver.def(
    "SetPostAdvanceCallback",
    static_cast<void (TransientSolver::*)(std::function<void()>)>(
      &TransientSolver::SetPostAdvanceCallback),
    R"(
    Register a callback that runs after each advance within :meth:`Execute`.

    :param callback: Function invoked after the solver advances a timestep. Pass None
        to clear.
    :type callback: Optional[Callable[[], None]]
    )");
  transient_solver.def(
    "SetPostAdvanceCallback",
    py::overload_cast<std::nullptr_t>(&TransientSolver::SetPostAdvanceCallback),
    "Clear the PostAdvance callback by passing None.");
  transient_solver.def(
    "IsConverged",
    &TransientSolver::IsConverged,
    R"(
    Return True if the transport solve of every time step since :meth:`Initialize` converged.

    A time step whose solve reaches its iteration limit, fails, or applies no convergence test
    (a single sweep, ``petsc_richardson`` with ``l_max_its = 1``) makes this False until the
    next :meth:`Initialize`, so checking it after each :meth:`Advance` identifies the first
    such step. Returns False before the first time step.
    )"
  );
  transient_solver.def(
    "ComputeBalanceTable",
    [](const TransientSolver& self)
    {
      return BalanceTableToDict(self.ComputeBalanceTable());
    },
    R"(
    Compute and return the global balance table using the solver's normalization.
    This is a collective operation and must be called on all ranks.

    Returns
    -------
    dict
        Dictionary with the following entries:

        - ``absorption_rate``:
          Global absorption rate, approximately ``integral sigma_a * phi dV`` summed over
          groups and the full domain.
        - ``production_rate``:
          Global volumetric production/source rate used by the solver,
          approximately ``integral Q dV`` summed over groups and the full domain.
        - ``inflow_rate``:
          Global incoming boundary contribution integrated over incoming
          angular flux on boundaries. On reflecting boundaries the outgoing partial
          current is counted as the incoming one, so reflecting faces add equal
          amounts to ``inflow_rate`` and ``outflow_rate`` and cancel in ``balance``.
        - ``outflow_rate``:
          Global outgoing boundary contribution accumulated from face outflow
          tallies, including reflecting boundaries (see ``inflow_rate``).
        - ``balance``:
          Rate balance,
          ``production_rate + inflow_rate - absorption_rate - outflow_rate``.
        - ``initial_inventory``:
          Total particle inventory at the start of the timestep, computed as
          ``integral (1 / v_g) * phi_old dV`` summed over groups and the full domain.
        - ``final_inventory``:
          Total particle inventory at the end of the timestep, computed as
          ``integral (1 / v_g) * phi_new dV`` summed over groups and the full domain.
        - ``predicted_inventory_change``:
          Inventory change predicted by the current timestep balance, computed as
          ``dt * balance``.
        - ``actual_inventory_change``:
          Measured change in total particle inventory over the timestep, computed as
          ``final_inventory - initial_inventory``.
        - ``inventory_residual``:
          Mismatch between the measured and predicted timestep inventory
          changes, computed as
          ``actual_inventory_change - predicted_inventory_change``.

    Notes
    -----
    This solver applies no extra normalization to the balance table.

    The transient inventory terms currently use the end-of-step rate balance to
    estimate the timestep inventory change.
    )"
  );
  slv.attr("BackwardEuler") = 1.0;
  slv.attr("CrankNicolson") = 0.5;
  // clang-format on
}

// Wrap non-linear k-eigen solver
void
WrapNLKEigen(py::module& slv)
{
  // clang-format off
  // non-linear k-eigen solver
  auto non_linear_k_eigen_solver = py::class_<NonLinearKEigenSolver, std::shared_ptr<NonLinearKEigenSolver>,
                                              Solver>(
    slv,
    "NonLinearKEigenSolver",
    R"(
    Non-linear k-eigenvalue solver.

    Wrapper of :cpp:class:`opensn::NonLinearKEigenSolver`.

    Supports one or more groupsets when the groupsets run without WGDSA/TGDSA.
    If WGDSA or TGDSA is enabled on any groupset, the nonlinear k-eigen solve
    must use a single groupset; otherwise ``Initialize()`` raises a ``ValueError``.

    The eigenvector is returned with an arbitrary scale that depends on the solver, the
    acceleration method, and the starting flux. Normalize absolute quantities (power, reaction
    rates, leakage, balance rates) explicitly, e.g. with ``power_normalization_target``.
    )"
  );
  non_linear_k_eigen_solver.def(
    py::init(
      [](py::kwargs& params)
      {
        return NonLinearKEigenSolver::Create(kwargs_to_param_block(params));
      }
        ),
    R"(
    Construct a non-linear k-eigenvalue solver.

    Parameters
    ----------
    problem: pyopensn.solver.DiscreteOrdinatesProblem
        Existing discrete ordinates problem instance.
        Multiple groupsets are supported when groupset WGDSA/TGDSA is disabled.
        If WGDSA or TGDSA is enabled on any groupset, the nonlinear k-eigen
        solve must use a single groupset.
    nl_abs_tol: float, default=1.0e-8
        Non-linear absolute tolerance.
    nl_rel_tol: float, default=1.0e-8
        Non-linear relative tolerance.
    nl_sol_tol: float, default=1.0e-50
        Non-linear solution tolerance.
    nl_max_its: int, default=50
        Non-linear algorithm maximum iterations.
    l_abs_tol: float, default=1.0e-8
        Linear absolute tolerance. The value used is ``min(l_abs_tol, nl_abs_tol)``, so
        that the linear solves can reduce the nonlinear residual below ``nl_abs_tol``.
    l_rel_tol: float, default=1.0e-8
        Linear relative tolerance.
    l_div_tol: float, default=1.0e6
        Linear divergence tolerance.
    l_max_its: int, default=50
        Linear algorithm maximum iterations.
    l_gmres_restart_intvl: int, default=30
        GMRES restart interval.
    l_gmres_breakdown_tol: float, default=0.1
        GMRES breakdown tolerance.
    reset_phi0: bool, default=True
        If true, reinitializes scalar fluxes to 1.0.
    num_initial_power_iterations: int, default=0
        Number of initial power iterations before the non-linear solve.
    compute_balance: bool, default=False
        If true, print the balance table, with production scaled by ``1/k_eff``, after
        :meth:`Execute` completes.

    Notes
    -----
    PETSc convergence failures and iteration limits are reported through the
    solver status and log output, consistent with the other transport solvers.
    Invalid residual states are reported to PETSc as function-domain
    errors so SNES can backtrack or terminate with the appropriate reason.
    )"
  );
  non_linear_k_eigen_solver.def(
    "GetEigenvalue",
    &NonLinearKEigenSolver::GetEigenvalue,
    R"(
    Return the current k-eigenvalue.
    )"
  );
  non_linear_k_eigen_solver.def(
    "GetNumSweeps",
    &NonLinearKEigenSolver::GetNumSweeps,
    R"(
    Return the cumulative number of transport sweeps performed by all groupsets, including
    those of the initial power iterations.

    Returns
    -------
    int
        Total number of transport sweeps.
    )"
  );
  non_linear_k_eigen_solver.def(
    "IsConverged",
    &NonLinearKEigenSolver::IsConverged,
    R"(
    Return True if the nonlinear solve of the last :meth:`Execute` converged.

    Returns False if the nonlinear solver stopped on its iteration limit or another
    divergence reason, and before the first :meth:`Execute`.
    )"
  );
  non_linear_k_eigen_solver.def(
    "ComputeBalanceTable",
    [](const NonLinearKEigenSolver& self)
    {
      return BalanceTableToDict(self.ComputeBalanceTable());
    },
    R"(
    Compute and return the global balance table using the solver's normalization.
    This is a collective operation and must be called on all ranks.

    Returns
    -------
    dict
        Dictionary with the following entries:

        - ``absorption_rate``:
          Global absorption rate, approximately ``integral sigma_a * phi dV`` summed over
          groups and the full domain.
        - ``production_rate``:
          Global volumetric production/source rate used by the solver,
          approximately ``integral Q dV`` summed over groups and the full domain.
        - ``inflow_rate``:
          Global incoming boundary contribution integrated over incoming
          angular flux on boundaries. On reflecting boundaries the outgoing partial
          current is counted as the incoming one, so reflecting faces add equal
          amounts to ``inflow_rate`` and ``outflow_rate`` and cancel in ``balance``.
        - ``outflow_rate``:
          Global outgoing boundary contribution accumulated from face outflow
          tallies, including reflecting boundaries (see ``inflow_rate``).
        - ``balance``:
          Rate balance,
          ``production_rate + inflow_rate - absorption_rate - outflow_rate``.

    Notes
    -----
    For k-eigenvalue balance reporting, this solver scales the production term by
    ``1 / k_eff`` before forming both ``production_rate`` and ``balance``.
    )"
  );
  // clang-format on
}

// Wrap power iteration solvers
void
WrapPIteration(py::module& slv)
{
  // clang-format off
  // power iteration k-eigen solver
  auto pi_k_eigen_solver = py::class_<PowerIterationKEigenSolver, std::shared_ptr<PowerIterationKEigenSolver>,
                                      Solver>(
    slv,
    "PowerIterationKEigenSolver",
    R"(
    Power iteration k-eigenvalue solver.

    Wrapper of :cpp:class:`opensn::PowerIterationKEigenSolver`.

    The eigenvector is returned with an arbitrary scale that depends on the solver, the
    acceleration method, and the starting flux. Normalize absolute quantities (power, reaction
    rates, leakage, balance rates) explicitly, e.g. with ``power_normalization_target``.
    )"
  );
  pi_k_eigen_solver.def(
    py::init(
      [](py::kwargs& params)
      {
        return PowerIterationKEigenSolver::Create(kwargs_to_param_block(params));
      }
    ),
    R"(
    Construct a power iteration k-eigen solver.

    Parameters
    ----------
    problem: pyopensn.solver.DiscreteOrdinatesProblem
        Existing DiscreteOrdinatesProblem instance.
    acceleration: pyopensn.solver.DiscreteOrdinatesKEigenAcceleration
        Optional DiscreteOrdinatesKEigenAcceleration instance for acceleration. It must be
        created for the same ``problem``. ``Initialize()`` checks the requirements of the
        acceleration method against the problem and raises a ``ValueError`` listing every
        violation.
    max_iters: int, default = 1000
        Maximum power iterations allowed.
    k_tol: float, default = 1.0e-10
        Tolerance on the estimated relative error of the k-eigenvalue. The iteration converges
        when the relative change of k between power iterations is below ``k_tol * (1 - rho)``
        in two consecutive iterations, where ``rho`` is the convergence rate estimated from the
        change of the production-normalized flux. Values below ``1.0e-12`` are treated as
        ``1.0e-12``.
    reset_phi0: bool, default=True
        If true, reinitializes scalar fluxes to 1.0.
    compute_balance: bool, default=False
        If true, print the balance table, with production scaled by ``1/k_eff``, after
        :meth:`Execute` completes.

    Notes
    -----
    With more than one groupset, each power iteration makes one across-groupset pass unless the
    problem sets ``max_ags_iterations``.

    If the problem was constructed with ``options={'read_restart_path': ...}``,
    restart data is read during :meth:`Initialize`. If it was constructed with
    ``options={'restart_writes_enabled': True, ...}``, timed restart dumps may
    be written during the outer iteration loop and a final restart dump is
    written when execution completes.
    )"
  );
  pi_k_eigen_solver.def(
    "GetEigenvalue",
    &PowerIterationKEigenSolver::GetEigenvalue,
    R"(
    Return the current k-eigenvalue.
    )"
  );
  pi_k_eigen_solver.def(
    "GetNumPowerIterations",
    &PowerIterationKEigenSolver::GetNumPowerIterations,
    R"(
    Return the number of completed power iterations.
    )"
  );
  pi_k_eigen_solver.def(
    "GetNumSweeps",
    &PowerIterationKEigenSolver::GetNumSweeps,
    R"(
    Return the total number of transport sweeps applied by all WGS solvers.
    )"
  );
  pi_k_eigen_solver.def(
    "IsConverged",
    &PowerIterationKEigenSolver::IsConverged,
    R"(
    Return True if the last :meth:`Execute` met the convergence criteria.

    Convergence requires the k-eigenvalue test on ``k_tol`` to pass and, when an
    acceleration method is attached, its own convergence check (for CMFD, the balance
    residual) to pass. Returns False if ``max_iters`` was reached first.
    )"
  );
  pi_k_eigen_solver.def(
    "ComputeBalanceTable",
    [](const PowerIterationKEigenSolver& self)
    {
      return BalanceTableToDict(self.ComputeBalanceTable());
    },
    R"(
    Compute and return the global balance table using the solver's normalization.
    This is a collective operation and must be called on all ranks.

    Returns
    -------
    dict
        Dictionary with the following entries:

        - ``absorption_rate``:
          Global absorption rate, approximately ``integral sigma_a * phi dV`` summed over
          groups and the full domain.
        - ``production_rate``:
          Global volumetric production/source rate used by the solver,
          approximately ``integral Q dV`` summed over groups and the full domain.
        - ``inflow_rate``:
          Global incoming boundary contribution integrated over incoming
          angular flux on boundaries. On reflecting boundaries the outgoing partial
          current is counted as the incoming one, so reflecting faces add equal
          amounts to ``inflow_rate`` and ``outflow_rate`` and cancel in ``balance``.
        - ``outflow_rate``:
          Global outgoing boundary contribution accumulated from face outflow
          tallies, including reflecting boundaries (see ``inflow_rate``).
        - ``balance``:
          Rate balance,
          ``production_rate + inflow_rate - absorption_rate - outflow_rate``.

    Notes
    -----
    For k-eigenvalue balance reporting, this solver scales the production term by
    ``1 / k_eff`` before forming both ``production_rate`` and ``balance``.
    )"
  );
  // clang-format on
}

// Wrap LBS solver
void
WrapDiscreteOrdinatesKEigenAcceleration(py::module& slv)
{
  // clang-format off
  // discrete ordinates k-eigen acceleration base
  auto acceleration = py::class_<DiscreteOrdinatesKEigenAcceleration,
                                 std::shared_ptr<DiscreteOrdinatesKEigenAcceleration>>(
    slv,
    "DiscreteOrdinatesKEigenAcceleration",
    R"(
    Base class for discrete ordinates k-eigenvalue acceleration methods.

    Wrapper of :cpp:class:`opensn::DiscreteOrdinatesKEigenAcceleration`.
    )"
  );
  // SCDSA acceleration
  auto scdsa_acceleration = py::class_<SCDSAAcceleration,
                                       std::shared_ptr<SCDSAAcceleration>,
                                       DiscreteOrdinatesKEigenAcceleration>(
    slv,
    "SCDSAAcceleration",
    R"(
    Construct an SCDSA accelerator for the power iteration k-eigenvalue solver.

    Wrapper of :cpp:class:`opensn::SCDSAAcceleration`.

    Requires a single groupset and, with the AAH sweep, ``l_max_its`` greater than one.
    ``sdm='pwlc'`` is not supported in 1D spherical geometry. The owning solver checks these
    requirements in ``Initialize()``. With ``l_max_its`` greater than one, each power iteration
    performs one sweep after the groupset solve to form the scalar-flux change, i.e. one sweep more
    than ``l_max_its``.
    )"
  );
  scdsa_acceleration.def(
    py::init(
      [](py::kwargs& params)
      {
        return SCDSAAcceleration::Create(kwargs_to_param_block(params));
      }
    ),
    R"(
    SCDSA acceleration for the power iteration k-eigenvalue solver.

    Parameters
    ----------
    problem: pyopensn.solver.DiscreteOrdinatesProblem
        Existing DiscreteOrdinatesProblem instance.
    l_abs_tol: float, default=1.0e-10
        Residual tolerance of the diffusion solves, relative to the norm of their right-hand
        side (the name is historical).
    max_iters: int, default=100
        Maximum allowable iterations.
    verbose: bool, default=False
        If true, enables verbose output.
    petsc_options: str, default="ssss"
        Additional PETSc options.
    pi_max_its: int, default=50
        Maximum allowable iterations for inner power iterations.
    pi_k_tol: float, default=1.0e-10
        k-eigenvalue tolerance for the inner power iterations.
    sdm: str, default="pwld"
        Spatial discretization method to use for the diffusion solver. Valid choices are:
            - 'pwld' : Piecewise Linear Discontinuous
            - 'pwlc' : Piecewise Linear Continuous
    )"
  );
  // SMM acceleration
  auto smm_acceleration = py::class_<SMMAcceleration,
                                     std::shared_ptr<SMMAcceleration>,
                                     DiscreteOrdinatesKEigenAcceleration>(
    slv,
    "SMMAcceleration",
    R"(
    Construct an SMM accelerator for the power iteration k-eigenvalue solver.

    Wrapper of :cpp:class:`opensn::SMMAcceleration`.

    Requires a single groupset, isotropic (P0) scattering, ``options.save_angular_flux=True``,
    only vacuum and reflecting boundaries (an isotropic boundary is rejected even with zero
    strength), and, with the AAH sweep, ``l_max_its`` greater than one. The owning solver checks
    these requirements in ``Initialize()``.
    )"
  );
  smm_acceleration.def(
    py::init(
      [](py::kwargs& params)
      {
        return SMMAcceleration::Create(kwargs_to_param_block(params));
      }
    ),
    R"(
    SMM acceleration for the power iteration k-eigenvalue solver.

    Warnings
    --------
       SMM acceleration is **experimental** and should be used with caution!
       SMM acceleration only supports problems with isotropic scattering.

    Parameters
    ----------
    problem: pyopensn.solver.DiscreteOrdinatesProblem
        Existing DiscreteOrdinatesProblem instance.
    l_abs_tol: float, default=1.0e-10
        Residual tolerance of the diffusion solves, relative to the norm of their right-hand
        side (the name is historical).
    max_iters: int, default=100
        Maximum allowable iterations.
    verbose: bool, default=False
        If true, enables verbose output.
    petsc_options: str, default="ssss"
        Additional PETSc options.
    pi_max_its: int, default=50
        Maximum allowable iterations for inner power iterations.
    pi_k_tol: float, default=1.0e-10
        k-eigenvalue tolerance for the inner power iterations.
    sdm: str, default="pwlc"
        Spatial discretization method to use for the diffusion solver. Valid choices are:
            - 'pwld' : Piecewise Linear Discontinuous
            - 'pwlc' : Piecewise Linear Continuous
    )"
  );
  // CMFD acceleration
  auto cmfd_acceleration = py::class_<CMFDAcceleration,
                                      std::shared_ptr<CMFDAcceleration>,
                                      DiscreteOrdinatesKEigenAcceleration>(
    slv,
    "CMFDAcceleration",
    R"(
    Construct a CMFD accelerator for the power iteration k-eigenvalue solver.

    Wrapper of :cpp:class:`opensn::CMFDAcceleration`.
    )"
  );
  cmfd_acceleration.def(
    py::init(
      [](py::kwargs& params)
      {
        return CMFDAcceleration::Create(kwargs_to_param_block(params));
      }
    ),
    R"(
    CMFD acceleration for the power iteration k-eigenvalue solver.

    With CMFD, each power iteration performs a configured number of high-order WGS
    transport update iterations, then applies a bounded low-order scalar-flux correction.
    The defaults use one WGS update iteration, the partial-current closure, fixed
    correction relaxation, and a transport-current balance gate before outer power
    iteration is allowed to converge.

    Most users should tune only the common controls: ``coarse_mesh``, ``current_closure``,
    ``aggregation_size``, ``group_aggregation_size``, ``relaxation``,
    ``update_wgs_max_its``, ``update_wgs_abs_tol``, and
    ``balance_residual_tolerance``. Parameters marked "developer/debug" are exposed for
    investigations and regression testing, but normally should be left at their defaults.

    Parameters
    ----------
    problem: pyopensn.solver.DiscreteOrdinatesProblem
        Existing DiscreteOrdinatesProblem instance.
    coarse_mesh: str, default="local_aggregation"
        Common option. Coarse-mesh construction method. Valid choices are:
            - 'identity' : one CMFD coarse cell per transport cell
            - 'local_aggregation' : connected same-block fine cells aggregated locally
            - 'global_aggregation' : connected same-block fine cells aggregated across MPI ranks

        ``"local_aggregation"`` is the conservative default. ``"global_aggregation"`` can
        reduce the coarse problem size on larger distributed meshes. Global aggregation is a
        logical CMFD coarse-space construction: it does not repartition the transport mesh or
        change fine-cell ownership. Each global coarse cell has one owning rank, while ranks
        owning member fine cells keep membership records for restriction and prolongation.
        Aggregates are connected by face adjacency, remain within one mesh block, and may be
        smaller than ``aggregation_size`` near boundaries or disconnected regions. ``"identity"``
        is primarily for debugging and method comparisons.
    current_closure: str, default="partial"
        Common option. CMFD face-current closure. Valid choices are:
            - 'partial' : partial-current CMFD (pCMFD); each one-sided partial current carries
              half of the diffusion coupling plus its own transport correction
            - 'net' : match the signed transport current across each coarse face

        ``"partial"`` is the default and is recommended for production use: it is typically
        stable with an unrelaxed correction on both optically thin and thick coarse cells,
        whereas the net-current closure can oscillate or diverge on optically thick coarse
        cells. Both closures share the same fixed point. ``"net"`` can converge in fewer
        iterations on some problems with optically thin coarse cells.
    aggregation_size: int, default=32
        Common option. Target number of fine cells per aggregated coarse cell for
        ``coarse_mesh="local_aggregation"`` or ``coarse_mesh="global_aggregation"``.
        Larger values reduce CMFD matrix size and
        setup/solve cost, but make the spatial correction less detailed. Smaller values
        increase the coarse-system cost but may improve robustness.
    group_aggregation_size: int, default=1
        Common option. Number of transport energy groups per CMFD coarse group, not the
        final number of coarse groups. A value of 1 preserves the full transport group
        structure in the low-order system. To target ``N`` total coarse groups, use
        ``(num_groups + N - 1) // N``.
    relaxation: float, default=1.0
        Common option. Strictly positive relaxation factor applied to the CMFD scalar-flux correction.
        This is the requested correction strength. The correction limiter may damp or
        skip an individual correction if the requested update would produce an invalid
        k-eigenvalue or non-finite flux, or would change the sign of the scalar flux in
        some coarse cell and coarse group. The default
        applies the full correction. A value below 1 is more conservative, but even an
        ideal correction then leaves about ``1 - relaxation`` of the error; reduce it when
        corrections are repeatedly damped or skipped.
    update_wgs_max_its: int, default=1
        Common option. Maximum WGS iterations used before each CMFD correction. The
        default performs one transport update iteration per power iteration; larger
        values make each transport update more accurate but more expensive.
    update_wgs_abs_tol: float, default=1.0e-12
        Common option. WGS absolute tolerance used before each CMFD correction. When
        ``update_wgs_max_its=1``, this tolerance is normally not the stopping criterion
        because only one WGS iteration is allowed. It matters when more WGS iterations
        are allowed.
    balance_residual_tolerance: float, default=1.0e-6
        Common option. Restricted transport-current balance residual tolerance required
        before CMFD permits outer power-iteration convergence. This prevents false
        convergence when ``k_eff_change`` is small but the CMFD-restricted transport
        balance is still inconsistent. This is a consistency guard, not an estimate of
        the k-eigenvalue error. For large 3D cases it may need to be looser than the
        outer ``k_tol``; choose it tight enough to prevent false convergence but not so
        tight that it forces unnecessary asymptotic balance iterations.
    l_abs_tol: float, default=1.0e-7
        Developer/debug. Absolute residual tolerance for each CMFD coarse linear solve.
        This affects the low-order linear solve, not the outer transport convergence.
    l_rel_tol: float, default=1.0e-8
        Developer/debug. Relative residual tolerance for each iterative CMFD coarse
        linear solve. The solve converges when either ``l_rel_tol`` or ``l_abs_tol`` is
        met. Unlike the absolute tolerance, it does not depend on the coarse flux
        normalization.
    max_iters: int, default=100
        Developer/debug. Maximum iterations for each CMFD coarse linear solve. Increase
        only if the coarse KSP solve is reaching its iteration limit. CMFD skips a
        correction when any coarse linear solve fails to converge within this limit. If
        the log reports ``correction = skipped (coarse_linear_solve_not_converged)``,
        use a direct coarse solve for modest coarse systems, increase this limit for
        iterative solves, or provide stronger ``petsc_options``.
    verbose: bool, default=False
        Developer/debug. If true, prints CMFD diagnostic and timing metrics. These are
        useful for regression tests and performance studies but can be noisy in production.
    petsc_options: str, default=""
        Developer/debug. Additional PETSc options for the CMFD coarse solver. Used only
        for solver experiments, especially with ``coarse_solver_policy="petsc_options"``.
    pi_max_its: int, default=50
        Developer/debug. Maximum inner power iterations for the CMFD coarse k solve (at least 1).
        Increase only if the coarse k solve is not converging enough to give useful
        corrections.
    pi_k_tol: float, default=1.0e-8
        Developer/debug. Tolerance on the relative change of the coarse k-eigenvalue between
        CMFD coarse power iterations. This is separate from the outer
        ``PowerIterationKEigenSolver`` k tolerance.
    pi_phi_tol: float, default=1.0e-6
        Developer/debug. Tolerance on the estimated coarse flux-shape error of the CMFD coarse
        power iteration: the relative L2 change of the production-normalized coarse flux divided
        by ``1 - rho``, with the convergence rate ``rho`` estimated from successive changes. The
        coarse power iteration stops when both ``pi_k_tol`` and ``pi_phi_tol`` are met (at least
        two iterations) or after ``pi_max_its`` iterations.
    correction_max_attempts: int, default=10
        Developer/debug. Maximum CMFD correction damping attempts before skipping the
        correction for the current transport update. Each failed attempt halves the
        damping. If the log reports ``correction = skipped (sign_change_guard)``,
        first reduce ``relaxation`` or use finer spatial/energy aggregation; changing
        this option is mainly for diagnostics.
    correction_min_damping: float, default=1.0e-4
        Developer/debug. Minimum CMFD correction damping factor considered during correction
        limiting. If no admissible correction is found above this damping, CMFD skips
        the correction for that transport update.
    inactive_iterations: int, default=0
        Developer/debug. Number of initial power iterations before applying CMFD
        corrections. These iterations retain the WGS settings configured on the problem.
        CMFD applies its transport update controls starting with the first active iteration.
    coarse_solver_policy: str, default="auto"
        Developer/debug. Coarse solver policy. Valid choices are:
            - 'auto' : PETSc preonly+LU below ``coarse_direct_solve_threshold``
              global unknowns, GMRES with hypre BoomerAMG otherwise
            - 'direct' : PETSc preonly+LU
            - 'iterative' : GMRES with hypre BoomerAMG
            - 'petsc_options' : allow ``petsc_options`` to override the PETSc KSP/PC setup

        CMFD corrections from unconverged coarse linear solves are always skipped.
        The skipped correction uses the unaccelerated transport update for that power
        iteration, including the raw transport k update so that power iteration can continue
        to move.
    coarse_direct_solve_threshold: int, default=20000
        Developer/debug. Maximum global CMFD unknown count for automatic direct coarse
        solves when ``coarse_solver_policy="auto"``. Larger values make ``"auto"`` use
        direct solves for larger coarse systems; choose values based on available host
        memory and acceptable factorization cost.
    )"
  );
  cmfd_acceleration.def(
    "GetNumAppliedCorrections",
    &CMFDAcceleration::GetNumAppliedCorrections,
    R"(
    Return the number of CMFD corrections applied during the last power-iteration solve.

    Inactive iterations are not counted. The count is reset at the start of each solve.
    )"
  );
  cmfd_acceleration.def(
    "GetNumSkippedCorrections",
    &CMFDAcceleration::GetNumSkippedCorrections,
    R"(
    Return the number of CMFD corrections skipped during the last power-iteration solve.

    A correction is skipped when a coarse linear solve does not converge or when every
    damping attempt is rejected; the power iteration then uses the unaccelerated transport
    update. Inactive iterations are not counted. The count is reset at the start of each
    solve.
    )"
  );
  // clang-format on
}

// Wrap the solver components of OpenSn
void
py_solver(py::module& pyopensn)
{
  py::module slv = pyopensn.def_submodule("solver", "Solver module.");
  WrapProblem(slv);
  WrapSolver(slv);
  WrapLBS(slv);
  WrapUncollidedSolver(slv);
  WrapSteadyState(slv);
  WrapTransient(slv);
  WrapNLKEigen(slv);
  WrapDiscreteOrdinatesKEigenAcceleration(slv);
  WrapPIteration(slv);
}

} // namespace opensn

Open-Sn / opensn / 42359882336

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous