AstroSimBench/tests/test_part4_plotting.py at main · CSeph/AstroSimBench · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
"""Exercise saved-data plotting and visible child-process diagnostics."""

from contextlib import redirect_stdout
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import patch
import hashlib
import io
import json
import os
import shutil
import sys
import unittest
import runpart4
from experiments.part4 import exp4e_zlk, validate_zlk_sampling
from experiments.part4.postprocess_part4 import primary_run,peak_data,run_analysis


class Part4PlottingTests(unittest.TestCase):
    """Use measured-data fixtures while prohibiting new integrations."""

    def copy_tables(self, destination, candidates):
        source=primary_run("4E_zlk")
        names=["reference_summary.csv"]
        names += [f"a10_i{angle}/reference_path.csv" for angle in (20,60,120)]
        for angle in (60,120):
            names += [f"a10_i{angle}/initial_conditions.json",f"a10_i{angle}/reference_ias15_tight.npz"]
        if candidates:
            names += ["measurements.csv","sampling_validation.json"]
            names += [f"a10_i{angle}/reference_sampling_refined.npz" for angle in (60,120)]
        for name in names:
            target=destination/name
            target.parent.mkdir(parents=True,exist_ok=True)
            shutil.copyfile(source/name,target)

    def test_plot_only_rebuilds_validated_figures_without_integration(self):
        with TemporaryDirectory() as directory:
            path=Path(directory)
            self.copy_tables(path,candidates=True)
            before={p:hashlib.sha256(p.read_bytes()).hexdigest() for p in path.rglob("*") if p.is_file()}
            with patch.object(exp4e_zlk,"run_experiment") as experiment, \
                 patch.object(validate_zlk_sampling,"rebound_reference") as reference, \
                 redirect_stdout(io.StringIO()):
                exp4e_zlk.main(["--plot-only",str(path)])
            experiment.assert_not_called()
            reference.assert_not_called()
            self.assertEqual(len(list(path.glob("*.png"))),4)
            self.assertEqual(len(list(path.glob("*.svg"))),4)
            self.assertTrue(all(p.read_bytes().startswith(b"\x89PNG\r\n\x1a\n") for p in path.glob("*.png")))
            self.assertTrue(all(hashlib.sha256(p.read_bytes()).hexdigest()==digest for p,digest in before.items()))
            self.assertFalse((path/"completion.json").exists())

    def test_reference_figures_do_not_require_finished_candidates(self):
        with TemporaryDirectory() as directory:
            path=Path(directory)
            self.copy_tables(path,candidates=False)
            with redirect_stdout(io.StringIO()):
                exp4e_zlk.redraw_run(path)
            self.assertEqual(len(list(path.glob("*.png"))),3)
            self.assertFalse((path/"03_secular_objective.png").exists())
            self.assertFalse((path/"completion.json").exists())

    def test_failed_child_output_is_visible_and_retained(self):
        with TemporaryDirectory() as directory:
            log=Path(directory)/"lesson.log"
            screen=io.StringIO()
            command=[sys.executable,"-u","-c","import sys; print('stage started'); print('test failure',file=sys.stderr); sys.exit(7)"]
            with redirect_stdout(screen):
                result=runpart4._run_logged(command,log_path=log,env=os.environ.copy())
            self.assertEqual(result.returncode,7)
            self.assertEqual(log.read_text(),screen.getvalue())
            self.assertIn("stage started",screen.getvalue())
            self.assertIn("test failure",screen.getvalue())

    def test_peak_view_uses_refined_samples_when_available(self):
        with TemporaryDirectory() as directory:
            path=Path(directory)
            self.copy_tables(path,candidates=True)
            summary,_=peak_data(path,allow_coarse=True)
            self.assertTrue(all(r['source_npz'].endswith('reference_sampling_refined.npz') for r in summary))
            self.assertTrue(all(abs(r['reference_cadence_over_T']-.625)<1e-12 for r in summary))
            for angle in (60,120):(path/f'a10_i{angle}'/'reference_sampling_refined.npz').unlink()
            summary,_=peak_data(path,allow_coarse=True)
            self.assertTrue(all(r['source_npz'].endswith('reference_ias15_tight.npz') for r in summary))
            self.assertTrue(all(abs(r['reference_cadence_over_T']-5.)<1e-12 for r in summary))
            with self.assertRaises(FileNotFoundError):peak_data(path)

    def test_analysis_keeps_outputs_with_lessons_and_reuses_d_table(self):
        with TemporaryDirectory() as directory:
            output=Path(directory)
            source=primary_run('4D_cost_to_target')
            target=output/'4D_cost_to_target'/source.name/'part4_cost_to_target.csv'
            target.parent.mkdir(parents=True)
            shutil.copyfile(source/'part4_cost_to_target.csv',target)
            before=target.read_bytes()
            batch=source.parent.parent
            with redirect_stdout(io.StringIO()):run_analysis(output,batch)
            self.assertEqual(target.read_bytes(),before)
            self.assertFalse((output/'companion').exists())
            self.assertEqual(len(list(output.glob('4*/analysis'))),4)
            self.assertTrue((output/'4A_corrector/analysis/4A_processed_vs_WH4.png').is_file())
            self.assertTrue((output/'4B_periapsis/analysis/4B_gain_vs_resolution.svg').is_file())
            e_run=primary_run('4E_zlk',batch)
            self.assertTrue((output/'4E_zlk'/e_run.name/'02_resolution_paths.png').is_file())
            self.assertEqual(json.loads((output/'analysis_completion.json').read_text())['new_integrations'],0)
            with self.assertRaises(FileExistsError):run_analysis(output,batch)


if __name__=="__main__":
    unittest.main()