Index

A

AbstractInversionStrategy (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.transport_operator_base) AbstractLinearOperator_sn (class in jax_sn.operators.base_linear_operator) AbstractTransportOperator (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.transport_operator_base) accepts() (AbstractInversionStrategy method), [1] (ComposedMultiGroupChainStrategy method), [1] (ComposedSingleGroupTransportBase method) (MultiGroupChainGatherStrategy method) (RequiresNonDifferentiableData method) add_axis_triad() (in module jax_sn.tools.plotly) additional_shape (MassBuilder attribute), [1] (MassOperator attribute), [1] all_connection_face_quadratures (BasixLagrangianSimplex property), [1] (Element property), [1] all_operators() (ComposedSingleGroupChainStrategy method), [1] (ComposedSingleGroupSourceStrategy method), [1] (ComposedSingleGroupStrategy method), [1] analyze_grad_memory() (in module jax_sn.tools.memory_analysis) analyze_grad_timing() (in module jax_sn.tools.memory_analysis)

analyze_memory() (OperatorAnalysis static method) analyze_timing() (OperatorAnalysis static method) angle (TransportBuilder attribute) (TransportData attribute), [1] angles (QuadratureSetBase attribute), [1] angles_original (QuadratureSetReduced attribute), [1] angular (DiscretisationData attribute) angular_data (OperatorData property) angular_data() (in module jax_sn.operator_creation.builders) AngularData (class in jax_sn.operators.operator_data) as_matrix() (AbstractLinearOperator_sn method) assert_compatible() (Element method), [1] assert_n_energy_groups() (CrossSectionData method), [1] assert_pytree_allclose() (in module jax_sn.tools) (in module jax_sn.tools.memory_analysis) assert_valid() (CrossSectionData method), [1] assert_valid_batched() (CrossSectionData method), [1] assume_full_rank() (DirectInverseSolver method) (MultiGroupLowerTriangularBlockSolver method), [1] available_quantities() (ParametrisedModule method), [1]

B

base_module (ParametrisedModuleTheta attribute) BaseOperatorData (class in jax_sn.operators.base_linear_operator) BasicSweepData (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.sweep_data_implementations) BasixLagrangianSimplex (class in jax_sn.solution_domain) (class in jax_sn.solution_domain.basix_fem) batch_size (BatchedForward attribute), [1] (BatchedPallasSweep attribute), [1] (BatchedSweep attribute), [1] BatchedForward (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.transport_operator_forward) BatchedPallasSweep (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.transport_operator_pallas)

BatchedSweep (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.transport_operator_jax) BroadcastedSingleGroupGather (class in jax_sn.operators.scatter_operator) (class in jax_sn.operators.scatter_operator.gather_scatter_operator_strategies) build_default_non_differentiable_data() (ComposedMultiGroupChainStrategy method), [1] (ComposedSingleGroupTransportBase method) (MultiGroupChainGatherStrategy method) (RequiresNonDifferentiableData method) build_default_sweep_data() (AbstractInversionStrategy method), [1] (ComposedMultiAngleInversionStrategy method), [1] (WarpLevelAngles method), [1] (WarpLevelSweep method), [1] build_kernel() (in module jax_sn.operators.mass_operator.mass_operator_warp) build_kernels() (in module jax_sn.operators.transport_operator.implementations.simplex_warp_kernel)

C

chain_strategy (MultiGroupLowerTriangularBlockSolver attribute), [1] check_compatible() (AbstractInversionStrategy method), [1] (ComposedMultiAngleInversionStrategy method), [1] (RequiresNonDifferentiableData method) (WarpLevelAngles method), [1] (WarpLevelSweep method), [1] (WarpSequentialAngles method), [1] (WarpSequentialSweep method), [1] compatible() (BasixLagrangianSimplex method), [1] (Element method), [1] ComposedGatherScatterStrategy (class in jax_sn.operators.scatter_operator) (class in jax_sn.operators.scatter_operator.gather_scatter_operator_base) ComposedMultiAngleForwardStrategy (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.transport_operator_strategies) ComposedMultiAngleInversionStrategy (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.transport_operator_strategies) ComposedMultiGroupChainStrategy (class in jax_sn.operators.multi_group_operator) (class in jax_sn.operators.multi_group_operator.composed_multi_group_strategies) ComposedSingleGroupChainStrategy (class in jax_sn.operators.single_group_operator) (class in jax_sn.operators.single_group_operator.composed_single_group_strategies) ComposedSingleGroupSourceStrategy (class in jax_sn.operators.single_group_operator) (class in jax_sn.operators.single_group_operator.composed_single_group_strategies) ComposedSingleGroupStrategy (class in jax_sn.operators.single_group_operator) (class in jax_sn.operators.single_group_operator.composed_single_group_strategies) ComposedSingleGroupTransportBase (class in jax_sn.operators.single_group_operator.composed_single_group_strategies) compute() (DirectInverseSolver method) (MultiGroupLowerTriangularBlockSolver method), [1] compute_3d_wk_yklm_real_from_quadrature_set() (in module jax_sn.quadrature_set) (in module jax_sn.quadrature_set.quadrature_set_base) compute_3d_yklm_real_from_quadrature_set() (in module jax_sn.quadrature_set) (in module jax_sn.quadrature_set.quadrature_set_base) compute_face_inflow() (in module jax_sn.domain.graph_generation) compute_face_matrices_linear_jacobian() (in module jax_sn.solution_domain.base_fem) compute_face_normal_vectors_from_3d_simplices() (in module jax_sn.tools.mesh_tools) compute_face_normals() (in module jax_sn.domain.geometric_domain) compute_geometric_data() (BasixLagrangianSimplex method), [1] (Element method), [1] compute_levels_domain() (in module jax_sn.domain.graph_generation) compute_levels_domain_safe() (in module jax_sn.domain.graph_generation) compute_levels_domain_safe_jax() (in module jax_sn.domain.graph_generation) compute_levels_iterative() (in module jax_sn.tools.graph_tools)

compute_levels_iterative_jax() (in module jax_sn.tools.graph_tools) compute_overlap_matrix() (in module jax_sn.energy_set) (in module jax_sn.energy_set.energy_set_overlap) compute_setup_for_solution_domain() (in module jax_sn.solution_domain) (in module jax_sn.solution_domain.solution_domain) compute_topology() (in module jax_sn.domain.geometric_domain) compute_topology_numpy() (in module jax_sn.tools.mesh_tools_numpy) compute_volume_matrices_linear_jacobian() (in module jax_sn.solution_domain.base_fem) conditionally_free (Graph attribute) conj() (AbstractLinearOperator_sn method) (DirectInverseSolver method) (MultiGroupLowerTriangularBlockSolver method), [1] connectivity (GeometricDomain attribute), [1] (SpatialBuilder attribute), [1] constant() (CrossSectionData class method), [1] constraints() (ParametrisedScipyInterface method) containing_cell() (GeometricDomain method), [1] contract() (EinsumFluxMoment method), [1] (FluxMomentStrategy method), [1] (StreamingFluxMoment method), [1] convert_to_downscattering_only() (CrossSectionData method), [1] create_all_combined_face_quadratures() (in module jax_sn.solution_domain.basix_fem) create_bvh() (in module jax_sn.tools.mesh_tools) create_bvh_cached() (in module jax_sn.tools.mesh_tools) create_cartesian_setup() (in module jax_sn.domain.domain_generation) create_face_centroids_from_simplices() (in module jax_sn.tools.mesh_tools) create_tn_quadrature_set() (in module jax_sn.quadrature_set) (in module jax_sn.quadrature_set.tn) cross_section_data (CrossSectionsIndexed attribute), [1] cross_sections (TransportData property), [1] cross_sections_indexed (Domain attribute), [1] (GatherScatterData attribute), [1] (ScatterData attribute), [1] (SpatialBuilder attribute), [1] CrossSectionData (class in jax_sn.domain) (class in jax_sn.domain.cross_sections) CrossSectionDataSingleGroup (class in jax_sn.domain) (class in jax_sn.domain.cross_sections) CrossSectionsIndexed (class in jax_sn.domain) (class in jax_sn.domain.cross_sections) custom_timeit() (in module jax_sn.tools.memory_analysis)

D

data (AbstractTransportOperator attribute), [1] (DataStrategyLinearOperator attribute) (DataTransposeDataStrategyLinearOperator attribute) (DiscreteOperator attribute), [1] (FluxMomentOperator attribute), [1] (GatherScatterOperator attribute), [1] (MassOperator attribute), [1] (MomentOperator attribute), [1] (MultiGroupOperatorBase attribute), [1] (ScatterOperator attribute), [1] (SingleGroupOperatorBase attribute) (TransportOperatorMultiAngle attribute), [1] (TransportOperatorSingleAngle attribute), [1] DataStrategyLinearOperator (class in jax_sn.operators.base_linear_operator) DataTransposeDataStrategyLinearOperator (class in jax_sn.operators.base_linear_operator) default_multi_group_pallas_strategy() (DefaultMultiGroupStrategy class method), [1] default_multi_group_strategy() (DefaultMultiGroupStrategy class method), [1] default_multi_group_warp_strategy() (DefaultMultiGroupStrategy class method), [1] DefaultMultiGroupStrategy (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.default_strategies) deps (Graph attribute) DERIVATIONS (DiscreteBuilder attribute), [1] (DiscretisationBuilder attribute), [1] (MassBuilder attribute), [1] (MomentBuilder attribute), [1] (MultiAngleTransportBuilder attribute), [1] (MultiGroupBuilder attribute), [1] (ParametrisedModule attribute), [1], [2], [3] (ScatterBuilder attribute), [1] (SingleAngleTransportBuilder attribute), [1] (SingleGroupBuilder attribute), [1] (SingleGroupBuilderBase attribute) (SingleGroupChainBuilder attribute), [1] (SingleGroupSourceBuilder attribute), [1] (SpatialBuilder attribute), [1] (SweptBuilder attribute), [1] (TransportBuilder attribute)

derivative_matrix (FEMMatrices attribute), [1] derived (Graph attribute) diagnostics (MultiGroupLowerTriangularBlockSolver attribute), [1] diagonal_strategy (MultiGroupLowerTriangularBlockSolver attribute), [1] DiagParametrisedScipyInterface (class in jax_sn.tools.parametrisation_tools) dij_f_p_fm_kernel_idx_in_precomp() (in module jax_sn.operators.transport_operator.implementations.transport_operator_jax) dij_f_p_fm_kernel_idx_in_precomp_vmap() (in module jax_sn.operators.transport_operator.implementations.transport_operator_jax) dimension (SolutionDomain property), [1] DirectInverseSolver (class in jax_sn.operators.base_linear_operator) DirectlyInvertibleOperator (class in jax_sn.operators.base_linear_operator) discrete_data() (in module jax_sn.operator_creation.builders) discrete_operator() (in module jax_sn.operator_creation.builders) discrete_strategy (ComposedSingleGroupChainStrategy attribute), [1] (ComposedSingleGroupStrategy attribute), [1] DiscreteBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) DiscreteData (class in jax_sn.operators.moment_operator) (class in jax_sn.operators.moment_operator.moment_discrete_operator_base) DiscreteMomentData (class in jax_sn.operators.moment_operator) (class in jax_sn.operators.moment_operator.moment_discrete_operator_base) DiscreteOperator (class in jax_sn.operators.moment_operator) (class in jax_sn.operators.moment_operator.moment_discrete_operator_base) discretisation_data (OperatorData attribute) discretisation_data() (in module jax_sn.operator_creation.builders) DiscretisationBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) DiscretisationData (class in jax_sn.operators.operator_data) Domain (class in jax_sn.domain) (class in jax_sn.domain.geometric_domain) domain (OperatorData property) (SolutionDomain property), [1] (SpatialData attribute), [1] domain() (in module jax_sn.operator_creation.builders) downscattering_only (CrossSectionData property), [1]

E

EinsumFluxMoment (class in jax_sn.operators.moment_operator) (class in jax_sn.operators.moment_operator.fluxmoment_strategies) Element (class in jax_sn.solution_domain) (class in jax_sn.solution_domain.base_fem) element (SolutionDomain property), [1] (SpatialBuilder attribute), [1] (SpatialData attribute), [1] elements_basis_shape (DiscreteOperator attribute), [1] (FluxMomentOperator attribute), [1] (MomentOperator attribute), [1] (OperatorData property) (SolutionDomain property), [1] (SpatialData property), [1] ends (LevelBasedParallelism attribute), [1] energy_group (IsotropicSourceSingleGroup attribute), [1] (SingleGroupBuilderBase attribute) (SingleGroupData attribute), [1] (TransportBuilder attribute) (TransportData attribute), [1]

energy_group_in (ScatterBuilder attribute), [1] (ScatterData attribute), [1] energy_group_out (ScatterBuilder attribute), [1] (ScatterData attribute), [1] etoe (GeometricDomain attribute), [1] (TopologyBuild attribute) etof (GeometricDomain attribute), [1] (TopologyBuild attribute) evaluate_jac() (DiagParametrisedScipyInterface method) (ParametrisedScipyInterface method) evaluate_values() (DiagParametrisedScipyInterface method) (ParametrisedScipyInterface method) extract_boundaries_from_mesh() (in module jax_sn.tools.mesh_tools) extract_elements_from_mesh() (in module jax_sn.tools.mesh_tools)

F

f_theta_parametrisation_values (ParametrisedScipyInterface attribute) face_jacobians_det (PrecomputedGeometricData attribute), [1] face_normals (GeometricDomain attribute), [1] face_normals() (in module jax_sn.operator_creation.builders) face_perm_idx (GeometricDomain attribute), [1] (TopologyBuild attribute) face_quadratures (BasixLagrangianSimplex property), [1] (Element property), [1] face_template (BasixLagrangianSimplex property), [1] (Element property), [1] (GeometricDomain attribute), [1] face_template() (in module jax_sn.operator_creation.builders) fem_matrices (MassData attribute), [1] (SolutionDomain property), [1] (SpatialData attribute), [1] fem_matrices() (in module jax_sn.operator_creation.builders) FEMMatrices (class in jax_sn.solution_domain) (class in jax_sn.solution_domain.base_fem) first_occurences (LevelBasedParallelism attribute), [1] FluxMomentOperator (class in jax_sn.operators.moment_operator) (class in jax_sn.operators.moment_operator.moment_discrete_operator_base) FluxMomentStrategy (class in jax_sn.operators.moment_operator) (class in jax_sn.operators.moment_operator.moment_discrete_operator_base) forward() (BatchedForward method), [1] (ComposedGatherScatterStrategy method), [1] (ComposedMultiGroupChainStrategy method), [1] (ComposedSingleGroupChainStrategy method), [1] (ComposedSingleGroupSourceStrategy method), [1] (ComposedSingleGroupStrategy method), [1] (FluxMomentStrategy method), [1] (IndexedScatter method), [1] (LaxMass method), [1] (MassStrategy method), [1] (MultiAngleForwardStrategy method), [1] (MultiGroupChainGatherStrategy method) (MultiGroupStrategyBase method), [1] (ScanAnglesForward method), [1] (SingleAngleForwardStrategy method), [1] (SingleGroupStrategyBase method), [1] (Strategy method) (VmapAnglesForward method), [1] (VmapForward method), [1] (VmapGatherScatterStrategy method), [1] (VmapMass method), [1] (WarpMass method), [1]

forward_args (TransportData property), [1] forward_strategy (MultiAngleTransportBuilder attribute), [1] (SingleAngleTransportBuilder attribute), [1] forward_transport_operator_batch() (in module jax_sn.operators.transport_operator.implementations.transport_operator_forward) forward_transport_operator_vmap() (in module jax_sn.operators.transport_operator.implementations.transport_operator_forward) ForwardDataStrategyLinearOperator (class in jax_sn.operators.base_linear_operator) free (Graph attribute) freezing() (ParametrisedModule class method), [1] from_basic_sweep_data() (LeveledSweepData class method), [1] from_composed() (ComposedMultiGroupChainStrategy class method), [1] (MultiGroupLowerTriangularBlockSolver class method), [1] from_domain_angle() (BasicSweepData class method), [1] (LeveledSweepData class method), [1] from_domain_quadrature_set() (BasicSweepData class method), [1] (LeveledSweepData class method), [1] from_element() (FEMMatrices class method), [1] from_element_and_domain() (SolutionDomain class method), [1] from_geometric_domain_and_cross_sections() (Domain class method), [1] from_list() (CrossSectionData class method), [1] from_mesh() (GeometricDomain class method), [1] from_mesh_and_cross_sections() (Domain class method), [1] from_mesh_and_cross_sections_numpy() (Domain class method), [1] from_mesh_numpy() (GeometricDomain class method), [1] from_operator() (IsotropicSourceSingleGroup class method), [1] from_operator_and_isotropic_source() (TotalSource class method), [1] from_quadrature_set() (AngularData class method) (DiscreteData class method), [1] (DiscreteMomentData class method), [1] (MomentData class method), [1] (QuadratureSetReduced class method), [1] from_scipy_interface() (DiagParametrisedScipyInterface class method) from_shape() (IsotropicSourceSingleGroup class method), [1] from_solution_domain() (SpatialBuilder class method), [1] from_sweep_data() (LevelBasedParallelism class method), [1] frontier() (ParametrisedModule class method), [1] frozen (Graph attribute) frozen() (ParametrisedModule method), [1] fw_cadis_scalar_flux() (in module jax_sn.variance_reduction) (in module jax_sn.variance_reduction.fw_cadis)

G

gather (MultiGroupLowerTriangularBlockSolver attribute), [1] gather_group() (BroadcastedSingleGroupGather method), [1] (SingleGroupGatherScatterStrategy method), [1] GatherScatterData (class in jax_sn.operators.scatter_operator) (class in jax_sn.operators.scatter_operator.gather_scatter_operator_base) GatherScatterOperator (class in jax_sn.operators.scatter_operator) (class in jax_sn.operators.scatter_operator.gather_scatter_operator_base) GatherScatterStrategy (class in jax_sn.operators.scatter_operator) (class in jax_sn.operators.scatter_operator.gather_scatter_operator_base) geometric_domain() (in module jax_sn.operator_creation.builders) GeometricDomain (class in jax_sn.domain) (class in jax_sn.domain.geometric_domain)

gij_inv_kernel_idx_in_precomp() (in module jax_sn.operators.transport_operator.implementations.transport_operator_jax) gij_inv_kernel_idx_in_precomp_vmap() (in module jax_sn.operators.transport_operator.implementations.transport_operator_jax) gij_kernel_idx_in_precomp() (in module jax_sn.operators.transport_operator.implementations.transport_operator_jax) gij_kernel_idx_in_precomp_vmap() (in module jax_sn.operators.transport_operator.implementations.transport_operator_jax) grad_objective() (ParametrisedScipyInterface method) grad_objective_and_constraints() (ParametrisedModuleObjective method) (ParametrisedObjective method) Graph (class in jax_sn.tools.parametrisation) graph() (ParametrisedModule class method), [1] group_moments_elements_basis_shape (GatherScatterOperator attribute), [1]

H

has_linear_jacobian (BasixLagrangianSimplex property), [1] (Element property), [1]

I

in_axis (DiscreteData attribute), [1] (DiscreteMomentData attribute), [1] (MomentData attribute), [1] in_structure() (AbstractLinearOperator_sn method) (ForwardDataStrategyLinearOperator method) (TransposeBoolDataStrategyLinearOperator method) incoming_gather() (SingleGroupGatherScatterStrategy method), [1] IndexedScatter (class in jax_sn.operators.scatter_operator) (class in jax_sn.operators.scatter_operator.scatter_operator_strategies) indices (CrossSectionsIndexed attribute), [1] INFLOW_EPSILON (in module jax_sn.operators.transport_operator.implementations.simplex_warp_kernel) inflow_matrices (FEMMatrices attribute), [1] init() (DirectInverseSolver method) (MultiGroupLowerTriangularBlockSolver method), [1] inner (ComposedMultiAngleForwardStrategy attribute), [1] (ComposedMultiAngleInversionStrategy attribute), [1] integration_coeff (FEMMatrices attribute), [1] interpolate() (SolutionDomain method), [1] interpolate_from_local_coordinate() (Element method), [1] interpolate_from_real_coordinate() (Element method), [1] inv() (DirectlyInvertibleOperator method)

inversion_args (TransportData property), [1] inversion_strategy (AbstractTransportOperator attribute), [1] (MultiAngleTransportBuilder attribute), [1] (SingleAngleTransportBuilder attribute), [1] (TransportOperatorMultiAngle attribute), [1] (TransportOperatorSingleAngle attribute), [1] invert() (AbstractInversionStrategy method), [1] (BatchedPallasSweep method), [1] (BatchedSweep method), [1] (MultiAngleInversionStrategy method), [1] (ScanAnglesInversion method), [1] (VmapAnglesInversion method), [1] (VmapPallasSweep method), [1] (VmapSweep method), [1] (WarpLevelAngles method), [1] (WarpLevelSweep method), [1] (WarpSequentialAngles method), [1] (WarpSequentialSweep method), [1] isotropic_value (IsotropicSourceSingleGroup attribute), [1] IsotropicSourceSingleGroup (class in jax_sn.solvers) (class in jax_sn.solvers.right_hand_side) iteration_i (DiagParametrisedScipyInterface attribute)

J

jac_constraints() (ParametrisedScipyInterface method) jax_sn module jax_sn.analytical_solutions module jax_sn.analytical_solutions.slab module jax_sn.domain module jax_sn.domain.cross_sections module jax_sn.domain.domain_generation module jax_sn.domain.geometric_domain module jax_sn.domain.graph_generation module jax_sn.energy_set module jax_sn.energy_set.energy_set_examples module jax_sn.energy_set.energy_set_overlap module jax_sn.operator_creation module jax_sn.operator_creation.builders module jax_sn.operator_creation.default_strategies module jax_sn.operators module jax_sn.operators.base_linear_operator module jax_sn.operators.mass_operator module jax_sn.operators.mass_operator.mass_operator_base module jax_sn.operators.mass_operator.mass_operator_strategies module jax_sn.operators.mass_operator.mass_operator_warp module jax_sn.operators.moment_operator module jax_sn.operators.moment_operator.fluxmoment_strategies module jax_sn.operators.moment_operator.moment_discrete_operator_base module jax_sn.operators.multi_group_operator module jax_sn.operators.multi_group_operator.composed_multi_group_strategies module jax_sn.operators.multi_group_operator.multi_group_operator_base module jax_sn.operators.operator_data module jax_sn.operators.scatter_operator module jax_sn.operators.scatter_operator.gather_scatter_operator_base module jax_sn.operators.scatter_operator.gather_scatter_operator_strategies module jax_sn.operators.scatter_operator.scatter_operator_base module jax_sn.operators.scatter_operator.scatter_operator_strategies module jax_sn.operators.single_group_operator module jax_sn.operators.single_group_operator.composed_single_group_strategies module jax_sn.operators.single_group_operator.single_group_operators_base module

jax_sn.operators.sweep_data module jax_sn.operators.transport_operator module jax_sn.operators.transport_operator.implementations module jax_sn.operators.transport_operator.implementations.multi_angle_inversions module jax_sn.operators.transport_operator.implementations.simplex_warp_kernel module jax_sn.operators.transport_operator.implementations.transport_operator_forward module jax_sn.operators.transport_operator.implementations.transport_operator_jax module jax_sn.operators.transport_operator.implementations.transport_operator_pallas module jax_sn.operators.transport_operator.implementations.transport_operator_warp module jax_sn.operators.transport_operator.sweep_data_implementations module jax_sn.operators.transport_operator.transport_data module jax_sn.operators.transport_operator.transport_operator_base module jax_sn.operators.transport_operator.transport_operator_strategies module jax_sn.quadrature_set module jax_sn.quadrature_set.quadrature_set_base module jax_sn.quadrature_set.tn module jax_sn.solution_domain module jax_sn.solution_domain.base_fem module jax_sn.solution_domain.basix_fem module jax_sn.solution_domain.solution_domain module jax_sn.solvers module jax_sn.solvers.base_solvers module jax_sn.solvers.right_hand_side module jax_sn.tools module jax_sn.tools.graph_tools module jax_sn.tools.memory_analysis module jax_sn.tools.mesh_tools module jax_sn.tools.mesh_tools_numpy module jax_sn.tools.parametrisation module jax_sn.tools.parametrisation_tools module jax_sn.tools.permutation_tools module jax_sn.tools.plotly module jax_sn.variance_reduction module jax_sn.variance_reduction.continuous_adjoint module jax_sn.variance_reduction.fw_cadis module

K

ks_agg() (in module jax_sn.tools.parametrisation_tools)

L

l_max (CrossSectionsIndexed property), [1] LaxMass (class in jax_sn.operators.mass_operator) (class in jax_sn.operators.mass_operator.mass_operator_strategies) level_fwd (LeveledSweepData attribute), [1] (ProvidesLevels attribute), [1] level_multi_angle (SimplexWarpKernels attribute) level_rev (LeveledSweepData attribute), [1] (ProvidesLevels attribute), [1] level_single_angle (SimplexWarpKernels attribute) LevelBasedParallelism (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.sweep_data_implementations)

LeveledSweepData (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.sweep_data_implementations) levels (BasicSweepData attribute), [1] (LeveledSweepData attribute), [1] (SweepData attribute), [1], [2] levels_argsort (BasicSweepData attribute), [1] (LeveledSweepData attribute), [1] (SweepData attribute), [1], [2] local_coordinates() (BasixLagrangianSimplex method), [1] (Element method), [1]

M

make_memory_parametrised_grad() (in module jax_sn.tools.memory_analysis) map_cross_sections_to_full_array() (SolutionDomain method), [1] map_over_angles() (in module jax_sn.operators.transport_operator.implementations.multi_angle_inversions) map_to_full_array() (CrossSectionsIndexed method), [1] mass_data() (in module jax_sn.operator_creation.builders) mass_matrix (FEMMatrices attribute), [1] mass_operator() (in module jax_sn.operator_creation.builders) mass_strategy (ComposedSingleGroupTransportBase attribute) MassBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) MassData (class in jax_sn.operators.mass_operator) (class in jax_sn.operators.mass_operator.mass_operator_base) MassOperator (class in jax_sn.operators.mass_operator) (class in jax_sn.operators.mass_operator.mass_operator_base) MassStrategy (class in jax_sn.operators.mass_operator) (class in jax_sn.operators.mass_operator.mass_operator_base) materialise() (ParametrisedModule method), [1] matrix (DiscreteData property), [1] (DiscreteMomentData property), [1] (MomentData property), [1] MAX_LENGTH (in module jax_sn.operators.transport_operator.implementations.simplex_warp_kernel) maximum_number_of_levels (WarpLevelAngles attribute), [1] (WarpLevelSweep attribute), [1] MissingQuantityError module jax_sn jax_sn.analytical_solutions jax_sn.analytical_solutions.slab jax_sn.domain jax_sn.domain.cross_sections jax_sn.domain.domain_generation jax_sn.domain.geometric_domain jax_sn.domain.graph_generation jax_sn.energy_set jax_sn.energy_set.energy_set_examples jax_sn.energy_set.energy_set_overlap jax_sn.operator_creation jax_sn.operator_creation.builders jax_sn.operator_creation.default_strategies jax_sn.operators jax_sn.operators.base_linear_operator jax_sn.operators.mass_operator jax_sn.operators.mass_operator.mass_operator_base jax_sn.operators.mass_operator.mass_operator_strategies jax_sn.operators.mass_operator.mass_operator_warp jax_sn.operators.moment_operator jax_sn.operators.moment_operator.fluxmoment_strategies jax_sn.operators.moment_operator.moment_discrete_operator_base jax_sn.operators.multi_group_operator jax_sn.operators.multi_group_operator.composed_multi_group_strategies jax_sn.operators.multi_group_operator.multi_group_operator_base jax_sn.operators.operator_data jax_sn.operators.scatter_operator jax_sn.operators.scatter_operator.gather_scatter_operator_base jax_sn.operators.scatter_operator.gather_scatter_operator_strategies jax_sn.operators.scatter_operator.scatter_operator_base jax_sn.operators.scatter_operator.scatter_operator_strategies jax_sn.operators.single_group_operator jax_sn.operators.single_group_operator.composed_single_group_strategies jax_sn.operators.single_group_operator.single_group_operators_base jax_sn.operators.sweep_data jax_sn.operators.transport_operator jax_sn.operators.transport_operator.implementations jax_sn.operators.transport_operator.implementations.multi_angle_inversions jax_sn.operators.transport_operator.implementations.simplex_warp_kernel jax_sn.operators.transport_operator.implementations.transport_operator_forward jax_sn.operators.transport_operator.implementations.transport_operator_jax jax_sn.operators.transport_operator.implementations.transport_operator_pallas jax_sn.operators.transport_operator.implementations.transport_operator_warp jax_sn.operators.transport_operator.sweep_data_implementations jax_sn.operators.transport_operator.transport_data jax_sn.operators.transport_operator.transport_operator_base jax_sn.operators.transport_operator.transport_operator_strategies jax_sn.quadrature_set jax_sn.quadrature_set.quadrature_set_base jax_sn.quadrature_set.tn jax_sn.solution_domain jax_sn.solution_domain.base_fem jax_sn.solution_domain.basix_fem jax_sn.solution_domain.solution_domain jax_sn.solvers jax_sn.solvers.base_solvers jax_sn.solvers.right_hand_side jax_sn.tools jax_sn.tools.graph_tools jax_sn.tools.memory_analysis jax_sn.tools.mesh_tools jax_sn.tools.mesh_tools_numpy jax_sn.tools.parametrisation jax_sn.tools.parametrisation_tools jax_sn.tools.permutation_tools jax_sn.tools.plotly jax_sn.variance_reduction jax_sn.variance_reduction.continuous_adjoint jax_sn.variance_reduction.fw_cadis

moment_data() (in module jax_sn.operator_creation.builders) moment_operator() (in module jax_sn.operator_creation.builders) moment_strategy (ComposedSingleGroupTransportBase attribute) MomentBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) MomentData (class in jax_sn.operators.moment_operator) (class in jax_sn.operators.moment_operator.moment_discrete_operator_base) MomentOperator (class in jax_sn.operators.moment_operator) (class in jax_sn.operators.moment_operator.moment_discrete_operator_base) moments_elements_basis_shape (ScatterOperator attribute), [1] multi_angle (SimplexWarpKernels attribute) multi_angle_transport_operator() (in module jax_sn.operator_creation.builders) multi_group_operator() (in module jax_sn.operator_creation.builders) MultiAngleForwardStrategy (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.transport_operator_base) MultiAngleInversionStrategy (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.transport_operator_base) MultiAngleTransportBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) MultiGroupBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) MultiGroupChainGatherScanStrategy (class in jax_sn.operators.multi_group_operator) (class in jax_sn.operators.multi_group_operator.composed_multi_group_strategies) MultiGroupChainGatherStrategy (class in jax_sn.operators.multi_group_operator.composed_multi_group_strategies) MultiGroupChainGatherVmapStrategy (class in jax_sn.operators.multi_group_operator) (class in jax_sn.operators.multi_group_operator.composed_multi_group_strategies) MultiGroupChainOperator (class in jax_sn.operators.multi_group_operator) (class in jax_sn.operators.multi_group_operator.multi_group_operator_base) MultiGroupChainScanStrategy (class in jax_sn.operators.multi_group_operator) (class in jax_sn.operators.multi_group_operator.composed_multi_group_strategies) MultiGroupChainStrategy (class in jax_sn.operators.multi_group_operator) (class in jax_sn.operators.multi_group_operator.multi_group_operator_base) MultiGroupChainVmapStrategy (class in jax_sn.operators.multi_group_operator) (class in jax_sn.operators.multi_group_operator.composed_multi_group_strategies) MultiGroupLowerTriangularBlockSolver (class in jax_sn.solvers) (class in jax_sn.solvers.base_solvers) MultiGroupOperator (class in jax_sn.operators.multi_group_operator) (class in jax_sn.operators.multi_group_operator.multi_group_operator_base) MultiGroupOperatorBase (class in jax_sn.operators.multi_group_operator) (class in jax_sn.operators.multi_group_operator.multi_group_operator_base) MultiGroupStrategy (class in jax_sn.operators.multi_group_operator) (class in jax_sn.operators.multi_group_operator.multi_group_operator_base) MultiGroupStrategyBase (class in jax_sn.operators.multi_group_operator) (class in jax_sn.operators.multi_group_operator.multi_group_operator_base) multiple_angle_level_bounds_safe() (in module jax_sn.tools.graph_tools) mv() (AbstractLinearOperator_sn method) (DataStrategyLinearOperator method) (ForwardDataStrategyLinearOperator method) (TransposeBoolDataStrategyLinearOperator method)

N

n_angles (AbstractTransportOperator property), [1] (DiscreteMomentData property), [1] (FluxMomentOperator property), [1] (OperatorData property) (QuadratureSetBase property), [1] (TransportOperatorMultiAngle property), [1] (TransportOperatorSingleAngle property), [1] n_basis (BasixLagrangianSimplex property), [1] (Element property), [1] (FEMMatrices property), [1] (OperatorData property) (SimplexWarpKernels attribute) (SolutionDomain property), [1] (SpatialData property), [1] n_elements (GeometricDomain property), [1] (PrecomputedGeometricData property), [1] (SolutionDomain property), [1] (SpatialData property), [1] n_energy_groups (CrossSectionsIndexed property), [1] (MultiGroupOperatorBase property), [1] (OperatorData property)

n_faces (BasixLagrangianSimplex property), [1] (Element property), [1] (SimplexWarpKernels attribute) n_faces_per_element (GeometricDomain property), [1] n_in_iteration (DiagParametrisedScipyInterface attribute) n_levels (LevelBasedParallelism property), [1] n_moments (DiscreteMomentData property), [1] (FluxMomentOperator property), [1] (OperatorData property) n_vertices_per_face() (BasixLagrangianSimplex method), [1] (Element method), [1] native_multi_angle() (SingleAngleForwardStrategy class method), [1] (SingleAngleInversionStrategy class method), [1] (WarpLevelSweep class method), [1] (WarpSequentialSweep class method), [1] NDIM (in module jax_sn.operators.transport_operator.implementations.simplex_warp_kernel) NON_DIFFERENTIABLE (ParametrisedModule attribute), [1], [2], [3] (SweptBuilder attribute), [1] (TransportBuilder attribute) non_differentiable_data (OperatorData attribute) non_differentiable_data() (in module jax_sn.operator_creation.builders) NonDifferentiableData (class in jax_sn.operators.operator_data)

O

obj_constraint() (in module jax_sn.tools.parametrisation_tools) objective() (ParametrisedScipyInterface method) objective_and_constraints() (ParametrisedModuleObjective method) (ParametrisedObjective method) operator_data (SingleGroupData attribute), [1] operator_data() (in module jax_sn.operator_creation.builders) OperatorAnalysis (class in jax_sn.tools.memory_analysis)

OperatorData (class in jax_sn.operators.operator_data) original_angles_indices (QuadratureSetReduced attribute), [1] out_structure() (AbstractLinearOperator_sn method) (ForwardDataStrategyLinearOperator method) (TransposeBoolDataStrategyLinearOperator method) outflow_matrices (FEMMatrices attribute), [1] outgoing_gather() (SingleGroupGatherScatterStrategy method), [1]

P

ParametrisationValues (class in jax_sn.tools.parametrisation_tools) parametrised_objective (ParametrisedScipyInterface attribute) ParametrisedModule (class in jax_sn.tools) (class in jax_sn.tools.parametrisation) ParametrisedModuleObjective (class in jax_sn.tools.parametrisation_tools) ParametrisedModuleTheta (class in jax_sn.tools.parametrisation_tools) ParametrisedObjective (class in jax_sn.tools.parametrisation_tools) ParametrisedScipyInterface (class in jax_sn.tools.parametrisation_tools) plot_1d_solution() (in module jax_sn.tools.plotly) plot_2d_solution() (in module jax_sn.tools.plotly) plot_derivation_graph() (in module jax_sn.tools.plotly) plot_domain() (in module jax_sn.tools.plotly) plot_element_basis_functions() (in module jax_sn.tools.plotly) plot_moments_bar() (in module jax_sn.tools.plotly) plot_moments_spheres() (in module jax_sn.tools.plotly) plot_quadrature_set() (in module jax_sn.tools.plotly) plot_spherical_field() (in module jax_sn.tools.plotly) precompute_geometric_data() (Element method), [1] precomputed_geometric_data (MassData attribute), [1] (SolutionDomain property), [1] (SpatialData attribute), [1] precomputed_geometric_data() (in module jax_sn.operator_creation.builders) PrecomputedGeometricData (class in jax_sn.solution_domain) (class in jax_sn.solution_domain.base_fem) primitive_deps (Graph attribute) primitives (Graph attribute)

print_single_group (SolverDiagnostics attribute), [1] print_sweep (SolverDiagnostics attribute), [1] probe_bvh() (in module jax_sn.tools.mesh_tools) projected (Graph attribute) PROJECTIONS (DiscreteBuilder attribute), [1] (DiscretisationBuilder attribute), [1] projections (Graph attribute) PROJECTIONS (MassBuilder attribute), [1] (MomentBuilder attribute), [1] (MultiAngleTransportBuilder attribute), [1] (MultiGroupBuilder attribute), [1] (ParametrisedModule attribute), [1], [2], [3] (ScatterBuilder attribute), [1] (SingleAngleTransportBuilder attribute), [1] (SingleGroupBuilder attribute), [1] (SingleGroupBuilderBase attribute) (SingleGroupChainBuilder attribute), [1] (SingleGroupSourceBuilder attribute), [1] (SpatialBuilder attribute), [1] (SweptBuilder attribute), [1] (TransportBuilder attribute) ProvidesLevels (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.sweep_data_implementations) psi_isotropic_box() (in module jax_sn.analytical_solutions) (in module jax_sn.analytical_solutions.slab) pytree_nbytes() (in module jax_sn.tools.memory_analysis)

Q

quadrature_set (AngularData attribute) (DiscretisationBuilder attribute), [1] QuadratureSet (class in jax_sn.quadrature_set) (class in jax_sn.quadrature_set.quadrature_set_base)

QuadratureSetBase (class in jax_sn.quadrature_set) (class in jax_sn.quadrature_set.quadrature_set_base) QuadratureSetReduced (class in jax_sn.quadrature_set) (class in jax_sn.quadrature_set.quadrature_set_base) quantities (Graph attribute)

R

real_coordinates() (BasixLagrangianSimplex method), [1] (Element method), [1] reconstruct() (BasicSweepData method), [1] (LeveledSweepData method), [1] (SweepData method), [1], [2] required_capabilities() (AbstractInversionStrategy method), [1] (ComposedMultiAngleInversionStrategy method), [1] (WarpLevelAngles method), [1] (WarpLevelSweep class method), [1] required_deps (Graph attribute)

RequiresNonDifferentiableData (class in jax_sn.operators.operator_data) RETAINED_UNDER (DiscretisationBuilder attribute), [1] (ParametrisedModule attribute), [1], [2], [3] (SweptBuilder attribute), [1] (TransportBuilder attribute) reverse (BasicSweepData property), [1] (LeveledSweepData property), [1] (SweepData property), [1], [2] RightHandSide (class in jax_sn.solvers) (class in jax_sn.solvers.right_hand_side)

S

save_dir (DiagParametrisedScipyInterface attribute) ScanAnglesForward (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.multi_angle_inversions) ScanAnglesInversion (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.multi_angle_inversions) scatter_data() (in module jax_sn.operator_creation.builders) scatter_operator() (in module jax_sn.operator_creation.builders) scatter_strategy (ComposedSingleGroupStrategy attribute), [1] ScatterBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) ScatterData (class in jax_sn.operators.scatter_operator) (class in jax_sn.operators.scatter_operator.scatter_operator_base) scattering (CrossSectionData attribute), [1] (CrossSectionDataSingleGroup attribute), [1] ScatterOperator (class in jax_sn.operators.scatter_operator) (class in jax_sn.operators.scatter_operator.scatter_operator_base) ScatterStrategy (class in jax_sn.operators.scatter_operator) (class in jax_sn.operators.scatter_operator.scatter_operator_base) shape (RightHandSide property), [1] SimplexWarpKernels (class in jax_sn.operators.transport_operator.implementations.simplex_warp_kernel) single_angle (SimplexWarpKernels attribute) single_angle() (TransportData method), [1] single_angle_strategy() (ComposedMultiAngleForwardStrategy method), [1] (ComposedMultiAngleInversionStrategy method), [1] (MultiAngleForwardStrategy method), [1] (MultiAngleInversionStrategy method), [1] (WarpLevelAngles method), [1] (WarpSequentialAngles method), [1] single_angle_transport_operator() (in module jax_sn.operator_creation.builders) single_group_angle_structure() (in module jax_sn.operators.single_group_operator) (in module jax_sn.operators.single_group_operator.single_group_operators_base) single_group_chain_operator() (in module jax_sn.operator_creation.builders) single_group_chain_strategy (ComposedMultiGroupChainStrategy attribute), [1] (MultiGroupChainGatherStrategy attribute) single_group_data() (in module jax_sn.operator_creation.builders) single_group_gather_scatter_strategy (MultiGroupChainGatherStrategy attribute) single_group_moment_structure() (in module jax_sn.operators.single_group_operator) (in module jax_sn.operators.single_group_operator.single_group_operators_base) single_group_operator() (in module jax_sn.operator_creation.builders) single_group_solver (MultiGroupLowerTriangularBlockSolver attribute), [1] single_group_source_operator() (in module jax_sn.operator_creation.builders) single_group_strategy (ComposedGatherScatterStrategy attribute), [1] SingleAngleForwardStrategy (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.transport_operator_base) SingleAngleInversionStrategy (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.transport_operator_base) SingleAngleTransportBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) SingleGroupBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) SingleGroupBuilderBase (class in jax_sn.operator_creation.builders) SingleGroupChainBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) SingleGroupChainOperator (class in jax_sn.operators.single_group_operator) (class in jax_sn.operators.single_group_operator.single_group_operators_base) SingleGroupChainStrategy (class in jax_sn.operators.single_group_operator) (class in jax_sn.operators.single_group_operator.single_group_operators_base) SingleGroupData (class in jax_sn.operators.single_group_operator) (class in jax_sn.operators.single_group_operator.single_group_operators_base) SingleGroupGatherScatterStrategy (class in jax_sn.operators.scatter_operator) (class in jax_sn.operators.scatter_operator.gather_scatter_operator_base) SingleGroupOperator (class in jax_sn.operators.single_group_operator) (class in jax_sn.operators.single_group_operator.single_group_operators_base) SingleGroupOperatorBase (class in jax_sn.operators.single_group_operator.single_group_operators_base) SingleGroupSourceBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) SingleGroupSourceOperator (class in jax_sn.operators.single_group_operator) (class in jax_sn.operators.single_group_operator.single_group_operators_base) SingleGroupSourceStrategy (class in jax_sn.operators.single_group_operator) (class in jax_sn.operators.single_group_operator.single_group_operators_base)

SingleGroupStrategy (class in jax_sn.operators.single_group_operator) (class in jax_sn.operators.single_group_operator.single_group_operators_base) SingleGroupStrategyBase (class in jax_sn.operators.single_group_operator) (class in jax_sn.operators.single_group_operator.single_group_operators_base) sizes (LevelBasedParallelism attribute), [1] solution_domain (MultiGroupOperatorBase property), [1] (OperatorData property) solution_domain() (in module jax_sn.operator_creation.builders) SolutionDomain (class in jax_sn.solution_domain) (class in jax_sn.solution_domain.solution_domain) solve_continuous_adjoint() (in module jax_sn.variance_reduction) (in module jax_sn.variance_reduction.continuous_adjoint) solve_multigroup_operator_rhs() (in module jax_sn.solvers) (in module jax_sn.solvers.right_hand_side) SolverDiagnostics (class in jax_sn.solvers) (class in jax_sn.solvers.base_solvers) spatial (DiscretisationData attribute) spatial_data (OperatorData property) (SolutionDomain attribute), [1] (TransportData attribute), [1] spatial_data() (in module jax_sn.operator_creation.builders) SpatialBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) SpatialData (class in jax_sn.solution_domain) (class in jax_sn.solution_domain.solution_domain) stored (Graph attribute) stored() (ParametrisedModule method), [1] storing() (ParametrisedModule class method), [1] strategy (AbstractTransportOperator attribute), [1] Strategy (class in jax_sn.operators.base_linear_operator) strategy (DataStrategyLinearOperator attribute) (DiscreteBuilder attribute), [1] (DiscreteOperator attribute), [1] (FluxMomentOperator attribute), [1] (GatherScatterOperator attribute), [1] (MassBuilder attribute), [1] (MassOperator attribute), [1] (MomentBuilder attribute), [1] (MomentOperator attribute), [1] (MultiGroupBuilder attribute), [1] (MultiGroupChainOperator attribute), [1] (MultiGroupOperator attribute), [1] (MultiGroupOperatorBase attribute), [1] (ScatterBuilder attribute), [1] (ScatterOperator attribute), [1] (SingleGroupBuilder attribute), [1] (SingleGroupChainBuilder attribute), [1] (SingleGroupChainOperator attribute), [1] (SingleGroupOperator attribute), [1] (SingleGroupOperatorBase attribute) (SingleGroupSourceBuilder attribute), [1] (SingleGroupSourceOperator attribute), [1] (TransportOperatorMultiAngle attribute), [1] (TransportOperatorSingleAngle attribute), [1] (TransposeBoolDataStrategyLinearOperator attribute) StreamingFluxMoment (class in jax_sn.operators.moment_operator) (class in jax_sn.operators.moment_operator.fluxmoment_strategies) structure (RightHandSide attribute), [1] SUPPORTED (in module jax_sn.operators.transport_operator.implementations.simplex_warp_kernel) sweep_data (NonDifferentiableData attribute) (OperatorData property) (TransportData attribute), [1] sweep_data() (in module jax_sn.operator_creation.builders) sweep_pl_update() (in module jax_sn.operators.transport_operator.implementations.transport_operator_pallas) SweepData (class in jax_sn.operators.sweep_data) (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.sweep_data_implementations) SweptBuilder (class in jax_sn.operator_creation) (class in jax_sn.operator_creation.builders) SymmetricDataStrategyLinearOperator (class in jax_sn.operators.base_linear_operator) SymmetricStrategy (class in jax_sn.operators.base_linear_operator)

T

tabulate_basis_functions_first_derivative() (BasixLagrangianSimplex method), [1] (Element method), [1] theta_parametrised_module (ParametrisedModuleObjective attribute) to_energy_group() (AbstractTransportOperator method), [1] (ScatterOperator method), [1] (SingleGroupOperatorBase method) to_module() (ParametrisedModuleObjective method) (ParametrisedModuleTheta method) to_single_angle() (TransportOperatorMultiAngle method), [1] topology_build() (in module jax_sn.operator_creation.builders) TopologyBuild (class in jax_sn.operator_creation.builders) total (CrossSectionData attribute), [1] (CrossSectionDataSingleGroup attribute), [1] total_mass_kernel_lax() (in module jax_sn.operators.mass_operator.mass_operator_strategies) total_scatter_kernel_xs_indexed() (in module jax_sn.operators.scatter_operator.scatter_operator_strategies) total_source (TotalSource attribute), [1] TotalSource (class in jax_sn.solvers) (class in jax_sn.solvers.right_hand_side) transport_data() (in module jax_sn.operator_creation.builders) transport_inverse_sweep_batched() (in module jax_sn.operators.transport_operator.implementations.transport_operator_jax) transport_inverse_sweep_batched_pallas() (in module jax_sn.operators.transport_operator.implementations.transport_operator_pallas) transport_inverse_sweep_complete_vmap() (in module jax_sn.operators.transport_operator.implementations.transport_operator_jax) transport_inverse_sweep_complete_vmap_pallas() (in module jax_sn.operators.transport_operator.implementations.transport_operator_pallas) transport_inverse_sweep_multiple_angles_batched_pallas() (in module jax_sn.operators.transport_operator.implementations.transport_operator_pallas) transport_inversion_strategy (ComposedSingleGroupTransportBase attribute) transport_strategy (ComposedSingleGroupTransportBase attribute) transport_sweep_data() (in module jax_sn.operator_creation.builders) TransportBuilder (class in jax_sn.operator_creation.builders) TransportData (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.transport_data) TransportOperatorMultiAngle (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.transport_operator_base)

TransportOperatorSingleAngle (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.transport_operator_base) TransposableOperatorData (class in jax_sn.operators.base_linear_operator) TransposableStrategy (class in jax_sn.operators.base_linear_operator) transpose() (AbstractLinearOperator_sn method) (ComposedGatherScatterStrategy method), [1] (ComposedMultiGroupChainStrategy method), [1] (ComposedSingleGroupChainStrategy method), [1] (ComposedSingleGroupSourceStrategy method), [1] (ComposedSingleGroupStrategy method), [1] (DataTransposeDataStrategyLinearOperator method) (DirectInverseSolver method) (FluxMomentStrategy method), [1] (MultiGroupChainGatherStrategy method) (MultiGroupLowerTriangularBlockSolver method), [1] (MultiGroupStrategyBase method), [1] (SingleGroupStrategyBase method), [1] (SymmetricDataStrategyLinearOperator method) (SymmetricStrategy method) (TransposableStrategy method) (TransposeBoolDataStrategyLinearOperator method) (VmapGatherScatterStrategy method), [1] transpose_bool (DiscreteOperator attribute), [1] (FluxMomentOperator attribute), [1] (GatherScatterOperator attribute), [1] (MomentOperator attribute), [1] (MultiGroupOperatorBase attribute), [1] (SingleGroupOperatorBase attribute) (TransposeBoolDataStrategyLinearOperator attribute) TransposeBoolDataStrategyLinearOperator (class in jax_sn.operators.base_linear_operator) transposed (TransportData property), [1] (TransposableOperatorData property)

U

updated() (ParametrisedModule method), [1]

V

varied() (ParametrisedModule method), [1] varying() (ParametrisedModule class method), [1] vertices (GeometricDomain attribute), [1] (SpatialBuilder attribute), [1] vitaminj() (in module jax_sn.energy_set.energy_set_examples) VmapAnglesForward (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.multi_angle_inversions) VmapAnglesInversion (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.multi_angle_inversions) VmapForward (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.transport_operator_forward) VmapGatherScatterStrategy (class in jax_sn.operators.scatter_operator) (class in jax_sn.operators.scatter_operator.gather_scatter_operator_strategies) VmapMass (class in jax_sn.operators.mass_operator) (class in jax_sn.operators.mass_operator.mass_operator_strategies) vmappable() (BatchedPallasSweep class method), [1] (BatchedSweep class method), [1] (SingleAngleInversionStrategy class method), [1] (VmapPallasSweep class method), [1] (VmapSweep class method), [1] (WarpLevelSweep class method), [1] (WarpSequentialSweep class method), [1]

VmapPallasSweep (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.transport_operator_pallas) VmapSweep (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.transport_operator_jax) vol_jacobians (PrecomputedGeometricData attribute), [1] vol_jacobians_det (PrecomputedGeometricData attribute), [1] vol_jacobians_inv (PrecomputedGeometricData attribute), [1] volume_quadrature (BasixLagrangianSimplex property), [1] (Element property), [1] volume_quadrature_order() (BasixLagrangianSimplex method), [1] (Element method), [1]

W

warp_mass_apply() (in module jax_sn.operators.mass_operator.mass_operator_warp) WarpLevelAngles (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.transport_operator_warp) WarpLevelSweep (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.transport_operator_warp) WarpMass (class in jax_sn.operators.mass_operator) (class in jax_sn.operators.mass_operator.mass_operator_warp) WarpSequentialAngles (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.transport_operator_warp)

WarpSequentialSweep (class in jax_sn.operators.transport_operator) (class in jax_sn.operators.transport_operator.implementations.transport_operator_warp) weights (QuadratureSetBase attribute), [1] weights_original (QuadratureSetReduced attribute), [1] wk_yklm (MomentData attribute), [1] wk_ylm (AngularData attribute) wk_ylm_real() (QuadratureSet method), [1] (QuadratureSetBase method), [1] (QuadratureSetReduced method), [1]

Y

yklm (DiscreteData attribute), [1] ylm (AngularData attribute)

ylm_real() (QuadratureSet method), [1] (QuadratureSetBase method), [1] (QuadratureSetReduced method), [1]

Z

zero() (CrossSectionData class method), [1]