Skip to content

Commit

Permalink
Match Jack's opacity state setup for easier fixing later
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewfullard committed Oct 7, 2024
1 parent a0fcb09 commit 9e3c968
Showing 1 changed file with 55 additions and 31 deletions.
86 changes: 55 additions & 31 deletions tardis/workflows/simple_tardis_workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,35 +36,30 @@ def __init__(self, configuration):
super().__init__(configuration, self.log_level, self.specific_log_level)
atom_data = parse_atom_data(configuration)

Check warning on line 37 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L35-L37

Added lines #L35 - L37 were not covered by tests

self.line_interaction_type = configuration.plasma.line_interaction_type

# set up states and solvers
self.simulation_state = SimulationState.from_config(

Check warning on line 40 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L40

Added line #L40 was not covered by tests
configuration,
atom_data=atom_data,
)

self.opacity_state = None
self.opacity_solver = OpacitySolver(
self.line_interaction_type,
configuration.plasma.disable_line_scattering,
)

self.macro_atom_state = None
if self.line_interaction_type in (
"downbranch",
"macroatom",
):
self.macro_atom_solver = MacroAtomSolver()
else:
self.macro_atom_solver = None

self.plasma_solver = assemble_plasma(

Check warning on line 45 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L45

Added line #L45 was not covered by tests
configuration,
self.simulation_state,
atom_data=atom_data,
)

line_interaction_type = configuration.plasma.line_interaction_type

Check warning on line 51 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L51

Added line #L51 was not covered by tests

self.opacity_solver = OpacitySolver(

Check warning on line 53 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L53

Added line #L53 was not covered by tests
line_interaction_type,
configuration.plasma.disable_line_scattering,
)

if line_interaction_type == "scatter":
self.macro_atom_solver = None

Check warning on line 59 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L58-L59

Added lines #L58 - L59 were not covered by tests
else:
self.macro_atom_solver = MacroAtomSolver()

Check warning on line 61 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L61

Added line #L61 was not covered by tests

self.transport_solver = MonteCarloTransportSolver.from_config(

Check warning on line 63 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L63

Added line #L63 was not covered by tests
configuration,
packet_source=self.simulation_state.packet_source,
Expand Down Expand Up @@ -326,11 +321,43 @@ def solve_plasma(self, estimated_radfield_properties):

self.plasma_solver.update(**update_properties)

Check warning on line 322 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L322

Added line #L322 was not covered by tests

def solve_montecarlo(self, no_of_real_packets, no_of_virtual_packets=0):
def solve_opacity(self):

Check warning on line 324 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L324

Added line #L324 was not covered by tests
"""Solves the opacity state and any associated objects
Returns
-------
dict
opacity_state : tardis.opacities.opacity_state.OpacityState
State of the line opacities
macro_atom_state : tardis.opacities.macro_atom.macro_atom_state.MacroAtomState or None
State of the macro atom
"""
opacity_state = self.opacity_solver.solve(self.plasma_solver)

Check warning on line 335 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L335

Added line #L335 was not covered by tests

if self.macro_atom_solver is None:
macro_atom_state = None

Check warning on line 338 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L337-L338

Added lines #L337 - L338 were not covered by tests
else:
macro_atom_state = self.macro_atom_solver.solve(

Check warning on line 340 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L340

Added line #L340 was not covered by tests
self.plasma_solver,
self.plasma_solver.atomic_data,
opacity_state.tau_sobolev,
self.plasma_solver.stimulated_emission_factor,
)

return {

Check warning on line 347 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L347

Added line #L347 was not covered by tests
"opacity_state": opacity_state,
"macro_atom_state": macro_atom_state,
}

def solve_montecarlo(

Check warning on line 352 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L352

Added line #L352 was not covered by tests
self, opacity_states, no_of_real_packets, no_of_virtual_packets=0
):
"""Solve the MonteCarlo process
Parameters
----------
opacity_states : dict
Opacity and (optionally) Macro Atom states.
no_of_real_packets : int
Number of real packets to simulate
no_of_virtual_packets : int, optional
Expand All @@ -343,10 +370,13 @@ def solve_montecarlo(self, no_of_real_packets, no_of_virtual_packets=0):
ndarray
Array of unnormalized virtual packet energies in each frequency bin
"""
opacity_state = opacity_states["opacity_state"]
macro_atom_state = opacity_states["macro_atom_state"]

Check warning on line 374 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L373-L374

Added lines #L373 - L374 were not covered by tests

transport_state = self.transport_solver.initialize_transport_state(

Check warning on line 376 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L376

Added line #L376 was not covered by tests
self.simulation_state,
self.opacity_state,
self.macro_atom_state,
opacity_state,
macro_atom_state,
self.plasma_solver,
no_of_real_packets,
no_of_virtual_packets=no_of_virtual_packets,
Expand Down Expand Up @@ -405,18 +435,10 @@ def run(self):
f"\n\tStarting iteration {(self.completed_iterations + 1):d} of {self.total_iterations:d}"
)

self.opacity_state = self.opacity_solver.solve(self.plasma_solver)

if self.macro_atom_solver is not None:
self.macro_atom_state = self.macro_atom_solver.solve(
self.plasma_solver,
self.plasma_solver.atomic_data,
self.opacity_state.tau_sobolev,
self.plasma_solver.stimulated_emission_factor,
)
opacity_states = self.solve_opacity()

Check warning on line 438 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L438

Added line #L438 was not covered by tests

transport_state, virtual_packet_energies = self.solve_montecarlo(

Check warning on line 440 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L440

Added line #L440 was not covered by tests
self.real_packet_count
opacity_states, self.real_packet_count
)

(

Check warning on line 444 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L444

Added line #L444 was not covered by tests
Expand All @@ -441,7 +463,9 @@ def run(self):
"\n\tITERATIONS HAVE NOT CONVERGED, starting final iteration"
)
transport_state, virtual_packet_energies = self.solve_montecarlo(

Check warning on line 465 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L465

Added line #L465 was not covered by tests
self.final_iteration_packet_count, self.virtual_packet_count
opacity_states,
self.final_iteration_packet_count,
self.virtual_packet_count,
)

self.initialize_spectrum_solver(

Check warning on line 471 in tardis/workflows/simple_tardis_workflow.py

View check run for this annotation

Codecov / codecov/patch

tardis/workflows/simple_tardis_workflow.py#L471

Added line #L471 was not covered by tests
Expand Down

0 comments on commit 9e3c968

Please sign in to comment.