#!/usr/bin/env python3 ############################################################################### # This file is part of SWIFT. # Copyright (c) 2023 Yolan Uyttenhove (yolan.uyttenhove@ugent.be) # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Lesser General Public License as published # by the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . # ############################################################################## from pathlib import Path import numpy as np import swiftsimio def test(name): def test_decorator(test_func): def test_runner(*args, **kwargs): try: test_func(*args, **kwargs) except Exception as e: print(f"\tTest {name} failed!") raise e else: print(f"\tTest {name} \u2705") return test_runner return test_decorator def approx_equal(a, b, threshold=0.001): return 2 * abs(a - b) / (abs(a) + abs(b)) < threshold @test("mass fractions sum to 1") def test_sum_mass_fractions(data): metal_fractions = data.gas.metal_mass_fractions h_fraction = data.gas.element_mass_fractions.hydrogen he_fraction = data.gas.element_mass_fractions.helium total = metal_fractions + he_fraction + h_fraction assert np.all(approx_equal(total, 1)) @test("total metal mass conservation") def test_total_metal_mass_conservation(data_start, data_end): def metal_mass(data): return data.gas.masses * data.gas.metal_mass_fractions assert approx_equal(np.sum(metal_mass(data_start)), np.sum(metal_mass(data_end))) def element_mass(data, element_name): return data.gas.masses * getattr(data.gas.element_mass_fractions, element_name) @test("hydrogen mass conservation") def test_h_mass_conservation(data_start, data_end): assert approx_equal( np.sum(element_mass(data_start, "hydrogen")), np.sum(element_mass(data_end, "hydrogen")), ) @test("helium mass conservation") def test_he_mass_conservation(data_start, data_end): assert approx_equal( np.sum(element_mass(data_start, "helium")), np.sum(element_mass(data_end, "helium")), ) @test("carbon mass conservation") def test_c_mass_conservation(data_start, data_end): assert approx_equal( np.sum(element_mass(data_start, "carbon")), np.sum(element_mass(data_end, "carbon")), ) @test("nitrogen mass conservation") def test_n_mass_conservation(data_start, data_end): assert approx_equal( np.sum(element_mass(data_start, "nitrogen")), np.sum(element_mass(data_end, "nitrogen")), ) @test("oxygen mass conservation") def test_o_mass_conservation(data_start, data_end): assert approx_equal( np.sum(element_mass(data_start, "oxygen")), np.sum(element_mass(data_end, "oxygen")), ) @test("neon mass conservation") def test_ne_mass_conservation(data_start, data_end): assert approx_equal( np.sum(element_mass(data_start, "neon")), np.sum(element_mass(data_end, "neon")) ) @test("magnesium mass conservation") def test_mg_mass_conservation(data_start, data_end): assert approx_equal( np.sum(element_mass(data_start, "magnesium")), np.sum(element_mass(data_end, "magnesium")), ) @test("silicon mass conservation") def test_si_mass_conservation(data_start, data_end): assert approx_equal( np.sum(element_mass(data_start, "silicon")), np.sum(element_mass(data_end, "silicon")), ) @test("iron mass conservation") def test_fe_mass_conservation(data_start, data_end): assert approx_equal( np.sum(element_mass(data_start, "iron")), np.sum(element_mass(data_end, "iron")) ) if __name__ == "__main__": print("Running sanity checks...") cwd = Path(__file__).parent start = swiftsimio.load(cwd / "output_0000.hdf5") end = swiftsimio.load(cwd / "output_0001.hdf5") test_sum_mass_fractions(end) test_total_metal_mass_conservation(start, end) test_h_mass_conservation(start, end) test_he_mass_conservation(start, end) test_c_mass_conservation(start, end) test_n_mass_conservation(start, end) test_o_mass_conservation(start, end) test_ne_mass_conservation(start, end) test_mg_mass_conservation(start, end) test_si_mass_conservation(start, end) test_fe_mass_conservation(start, end) print("Done!")