From d381f02450ec0aa5214857677bf71d05f7294d0a Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 18:56:21 -0400 Subject: [PATCH] fix shell info widget --- .../plasma/construction_simple_plasma.ipynb | 122 ++++++++++++------ tardis/visualization/widgets/shell_info.py | 2 +- .../widgets/tests/test_shell_info.py | 8 +- 3 files changed, 87 insertions(+), 45 deletions(-) diff --git a/docs/physics/plasma/construction_simple_plasma.ipynb b/docs/physics/plasma/construction_simple_plasma.ipynb index 00e94706acf..e0315f3f4c7 100644 --- a/docs/physics/plasma/construction_simple_plasma.ipynb +++ b/docs/physics/plasma/construction_simple_plasma.ipynb @@ -9,48 +9,12 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": 9, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/wkerzend/python/tardis/tardis/__init__.py:20: UserWarning: Astropy is already imported externally. Astropy should be imported after TARDIS.\n", - " warnings.warn(\n" - ] - }, - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "a4d91ae0803f463c96fe296874d0bacd", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Iterations: 0/? [00:00 1\u001b[0m \u001b[43mplasma_solver_factory\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43massemble\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnumber_densities\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43md_radfield\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m5\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mu\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mday\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/python/tardis/tardis/plasma/assembly/base.py:506\u001b[0m, in \u001b[0;36mPlasmaSolverFactory.assemble\u001b[0;34m(self, number_densities, dilute_planckian_radiation_field, time_explosion, electron_densities)\u001b[0m\n\u001b[1;32m 504\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msetup_electron_densities(electron_densities)\n\u001b[1;32m 505\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhelium_treatment\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhelium_treatment\n\u001b[0;32m--> 506\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mBasePlasma\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 507\u001b[0m \u001b[43m \u001b[49m\u001b[43mplasma_properties\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mplasma_modules\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 508\u001b[0m \u001b[43m \u001b[49m\u001b[43mproperty_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mproperty_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 509\u001b[0m \u001b[43m \u001b[49m\u001b[43mplasma_solver_settings\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mplasma_solver_settings\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 510\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 511\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/python/tardis/tardis/io/util.py:195\u001b[0m, in \u001b[0;36mHDFWriterMixin.__new__\u001b[0;34m(cls, *args, **kwargs)\u001b[0m\n\u001b[1;32m 193\u001b[0m instance \u001b[38;5;241m=\u001b[39m \u001b[38;5;28msuper\u001b[39m(HDFWriterMixin, \u001b[38;5;28mcls\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__new__\u001b[39m(\u001b[38;5;28mcls\u001b[39m)\n\u001b[1;32m 194\u001b[0m instance\u001b[38;5;241m.\u001b[39moptional_hdf_properties \u001b[38;5;241m=\u001b[39m []\n\u001b[0;32m--> 195\u001b[0m \u001b[43minstance\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 196\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m instance\n", + "File \u001b[0;32m~/python/tardis/tardis/plasma/base.py:39\u001b[0m, in \u001b[0;36mBasePlasma.__init__\u001b[0;34m(self, plasma_properties, property_kwargs, plasma_solver_settings, **kwargs)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mplasma_properties \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_init_properties(\n\u001b[1;32m 36\u001b[0m plasma_properties, property_kwargs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m 37\u001b[0m )\n\u001b[1;32m 38\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mplasma_solver_settings \u001b[38;5;241m=\u001b[39m plasma_solver_settings\n\u001b[0;32m---> 39\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_build_graph\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 40\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mupdate(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n", + "File \u001b[0;32m~/python/tardis/tardis/plasma/base.py:101\u001b[0m, in \u001b[0;36mBasePlasma._build_graph\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 99\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m \u001b[38;5;28minput\u001b[39m \u001b[38;5;129;01min\u001b[39;00m plasma_property\u001b[38;5;241m.\u001b[39minputs:\n\u001b[1;32m 100\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28minput\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutputs_dict:\n\u001b[0;32m--> 101\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m PlasmaMissingModule(\n\u001b[1;32m 102\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mModule \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mplasma_property\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m requires input \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 103\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28minput\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m which has not been added\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 104\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m to this plasma\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 105\u001b[0m )\n\u001b[1;32m 106\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 107\u001b[0m position \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutputs_dict[\u001b[38;5;28minput\u001b[39m]\u001b[38;5;241m.\u001b[39moutputs\u001b[38;5;241m.\u001b[39mindex(\u001b[38;5;28minput\u001b[39m)\n", + "\u001b[0;31mPlasmaMissingModule\u001b[0m: Module PartitionFunction requires input level_boltzmann_factor which has not been added to this plasma" + ] + } + ], + "source": [ + "plasma_solver_factory.assemble(number_densities, d_radfield, 5 * u.day)" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[tardis.plasma.properties.plasma_input.DilutePlanckianRadField,\n", + " tardis.plasma.properties.plasma_input.DilutePlanckianRadField,\n", + " tardis.plasma.properties.plasma_input.NumberDensity,\n", + " tardis.plasma.properties.plasma_input.TimeExplosion,\n", + " tardis.plasma.properties.plasma_input.AtomicData,\n", + " tardis.plasma.properties.plasma_input.JBlues,\n", + " tardis.plasma.properties.plasma_input.LinkTRadTElectron,\n", + " tardis.plasma.properties.plasma_input.HeliumTreatment,\n", + " tardis.plasma.properties.plasma_input.ContinuumInteractionSpecies,\n", + " tardis.plasma.properties.plasma_input.NLTEIonizationSpecies,\n", + " tardis.plasma.properties.plasma_input.NLTEExcitationSpecies,\n", + " tardis.plasma.properties.plasma_input.TRadiative,\n", + " tardis.plasma.properties.plasma_input.DilutionFactor,\n", + " tardis.plasma.properties.general.BetaRadiation,\n", + " tardis.plasma.properties.atomic.Levels,\n", + " tardis.plasma.properties.atomic.Lines,\n", + " tardis.plasma.properties.partition_function.PartitionFunction,\n", + " tardis.plasma.properties.general.GElectron,\n", + " tardis.plasma.properties.atomic.IonizationData,\n", + " tardis.plasma.properties.atomic.LinesLowerLevelIndex,\n", + " tardis.plasma.properties.atomic.LinesUpperLevelIndex,\n", + " tardis.opacities.tau_sobolev.TauSobolev,\n", + " tardis.plasma.properties.radiative_properties.StimulatedEmissionFactor,\n", + " tardis.plasma.properties.general.SelectedAtoms,\n", + " tardis.plasma.properties.general.ElectronTemperature,\n", + " tardis.plasma.properties.partition_function.LevelBoltzmannFactorLTE,\n", + " tardis.plasma.properties.ion_population.PhiSahaLTE,\n", + " tardis.plasma.properties.level_population.LevelNumberDensity,\n", + " tardis.plasma.properties.ion_population.IonNumberDensity]" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "plasma_solver_factory.plasma_modules" ] }, { diff --git a/tardis/visualization/widgets/shell_info.py b/tardis/visualization/widgets/shell_info.py index 811386be4ae..2cb93c91dd6 100644 --- a/tardis/visualization/widgets/shell_info.py +++ b/tardis/visualization/widgets/shell_info.py @@ -190,7 +190,7 @@ def __init__(self, sim_model): super().__init__( sim_model.simulation_state.t_radiative, sim_model.simulation_state.dilution_factor, - sim_model.plasma.abundance, + sim_model.simulation_state.abundance, sim_model.plasma.number_density, sim_model.plasma.ion_number_density, sim_model.plasma.level_number_density, diff --git a/tardis/visualization/widgets/tests/test_shell_info.py b/tardis/visualization/widgets/tests/test_shell_info.py index 7a99374dbc5..df7f4f168d9 100644 --- a/tardis/visualization/widgets/tests/test_shell_info.py +++ b/tardis/visualization/widgets/tests/test_shell_info.py @@ -16,7 +16,7 @@ def base_shell_info(simulation_verysimple): return BaseShellInfo( simulation_verysimple.simulation_state.t_radiative, simulation_verysimple.simulation_state.dilution_factor, - simulation_verysimple.plasma.abundance, + simulation_verysimple.simulation_state.abundance, simulation_verysimple.plasma.number_density, simulation_verysimple.plasma.ion_number_density, simulation_verysimple.plasma.level_number_density, @@ -58,12 +58,14 @@ def test_element_count_data( ): element_count_data = base_shell_info.element_count(1) assert element_count_data.shape == ( - len(simulation_verysimple.plasma.abundance[shell_num - 1]), + len( + simulation_verysimple.simulation_state.abundance[shell_num - 1] + ), 2, ) assert np.allclose( element_count_data.iloc[:, -1].map(np.float64), - simulation_verysimple.plasma.abundance[shell_num - 1], + simulation_verysimple.simulation_state.abundance[shell_num - 1], ) @pytest.mark.parametrize(("atomic_num", "shell_num"), [(12, 1), (20, 20)])