b2-run-sim 0.0.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- b2_run_sim-0.0.1/.flake8 +8 -0
- b2_run_sim-0.0.1/.gitignore +5 -0
- b2_run_sim-0.0.1/LICENSE +19 -0
- b2_run_sim-0.0.1/PKG-INFO +33 -0
- b2_run_sim-0.0.1/README.md +19 -0
- b2_run_sim-0.0.1/example/dec.dec +26 -0
- b2_run_sim-0.0.1/example/setup.bash +11 -0
- b2_run_sim-0.0.1/example/steer_recon.py +126 -0
- b2_run_sim-0.0.1/example/steer_sim.py +53 -0
- b2_run_sim-0.0.1/example.py +72 -0
- b2_run_sim-0.0.1/pyproject.toml +23 -0
- b2_run_sim-0.0.1/src/b2_run_sim/__init__.py +4 -0
- b2_run_sim-0.0.1/src/b2_run_sim/config.py +142 -0
- b2_run_sim-0.0.1/src/b2_run_sim/setup.py +117 -0
- b2_run_sim-0.0.1/src/b2_run_sim/submit.py +202 -0
- b2_run_sim-0.0.1/src/b2_run_sim/util.py +121 -0
b2_run_sim-0.0.1/.flake8
ADDED
b2_run_sim-0.0.1/LICENSE
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
Copyright (c) Ethan Lee
|
|
2
|
+
|
|
3
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
4
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
5
|
+
in the Software without restriction, including without limitation the rights
|
|
6
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
7
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
8
|
+
furnished to do so, subject to the following conditions:
|
|
9
|
+
|
|
10
|
+
The above copyright notice and this permission notice shall be included in all
|
|
11
|
+
copies or substantial portions of the Software.
|
|
12
|
+
|
|
13
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
14
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
15
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
16
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
17
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
18
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
19
|
+
SOFTWARE.
|
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: b2_run_sim
|
|
3
|
+
Version: 0.0.1
|
|
4
|
+
Summary: Simulation helper for basf2.
|
|
5
|
+
Project-URL: Homepage, https://github.com/ethanlee20/b2_sim
|
|
6
|
+
Project-URL: Issues, https://github.com/ethanlee20/b2_sim/issues
|
|
7
|
+
Author-email: Ethan Lee <elee20@hawaii.edu>
|
|
8
|
+
License-Expression: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Classifier: Operating System :: OS Independent
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Requires-Python: >=3.14
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
|
|
15
|
+
After setting up your environment, run using
|
|
16
|
+
|
|
17
|
+
```
|
|
18
|
+
python3.14 <script>
|
|
19
|
+
```
|
|
20
|
+
|
|
21
|
+
Steering files need to take the following command line arguments
|
|
22
|
+
|
|
23
|
+
```
|
|
24
|
+
simulation_steering_file.py <decay file path> <output file path> <number of events>
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
```
|
|
28
|
+
reconstruction_steering_file.py <simulated events file path> <output file path>
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
Decay file should be specified as a template. i.e. put `{parameter_name}` in place of parameter values.
|
|
32
|
+
|
|
33
|
+
See `example.py` script and `example` directory for more details.
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
After setting up your environment, run using
|
|
2
|
+
|
|
3
|
+
```
|
|
4
|
+
python3.14 <script>
|
|
5
|
+
```
|
|
6
|
+
|
|
7
|
+
Steering files need to take the following command line arguments
|
|
8
|
+
|
|
9
|
+
```
|
|
10
|
+
simulation_steering_file.py <decay file path> <output file path> <number of events>
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
```
|
|
14
|
+
reconstruction_steering_file.py <simulated events file path> <output file path>
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
Decay file should be specified as a template. i.e. put `{parameter_name}` in place of parameter values.
|
|
18
|
+
|
|
19
|
+
See `example.py` script and `example` directory for more details.
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
Alias MyB0 B0
|
|
2
|
+
Alias MyAntiB0 anti-B0
|
|
3
|
+
ChargeConj MyB0 MyAntiB0
|
|
4
|
+
|
|
5
|
+
Alias MyK*0 K*0
|
|
6
|
+
Alias MyAnti-K*0 anti-K*0
|
|
7
|
+
ChargeConj MyK*0 MyAnti-K*0
|
|
8
|
+
|
|
9
|
+
Decay Upsilon(4S)
|
|
10
|
+
0.500 MyB0 anti-B0 VSS;
|
|
11
|
+
0.500 B0 MyAntiB0 VSS;
|
|
12
|
+
Enddecay
|
|
13
|
+
|
|
14
|
+
Decay MyB0
|
|
15
|
+
1.000 MyK*0 mu+ mu- BTOSLLNPR 0 0 {dC_7} 0 1 {dC_9} 0 2 {dC_10} 0;
|
|
16
|
+
Enddecay
|
|
17
|
+
|
|
18
|
+
CDecay MyAntiB0
|
|
19
|
+
|
|
20
|
+
Decay MyK*0
|
|
21
|
+
1.000 K+ pi- VSS;
|
|
22
|
+
Enddecay
|
|
23
|
+
|
|
24
|
+
CDecay MyAnti-K*0
|
|
25
|
+
|
|
26
|
+
End
|
|
@@ -0,0 +1,126 @@
|
|
|
1
|
+
import sys
|
|
2
|
+
|
|
3
|
+
import basf2 as b2
|
|
4
|
+
import modularAnalysis as ma
|
|
5
|
+
from variables import collections as vc
|
|
6
|
+
from variables import utils as vu
|
|
7
|
+
from variables import variables as vm
|
|
8
|
+
|
|
9
|
+
# command line options
|
|
10
|
+
|
|
11
|
+
input_file_paths = sys.argv[1] # overwritten on the Grid
|
|
12
|
+
output_file_path = sys.argv[2] # choose something like recon.root
|
|
13
|
+
|
|
14
|
+
print("\n-- Reconstruction Configuration --")
|
|
15
|
+
print(f"Input file (overwritten on grid): {input_file_paths}")
|
|
16
|
+
print(f"Output file (modified on grid): {output_file_path}")
|
|
17
|
+
print("-----------------------------------\n")
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
# begin processing
|
|
21
|
+
|
|
22
|
+
main = b2.Path()
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def append_global_tag():
|
|
26
|
+
gt = ma.getAnalysisGlobaltag()
|
|
27
|
+
b2.conditions.append_globaltag(gt)
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def define_aliases():
|
|
31
|
+
vm.addAlias("mcMother_mcPDG", "mcMother(mcPDG)")
|
|
32
|
+
vm.addAlias("mcSister_0_mcPDG", "mcMother(mcDaughter(0, mcPDG))")
|
|
33
|
+
vm.addAlias("mcSister_1_mcPDG", "mcMother(mcDaughter(1, mcPDG))")
|
|
34
|
+
vm.addAlias("mcSister_2_mcPDG", "mcMother(mcDaughter(2, mcPDG))")
|
|
35
|
+
vm.addAlias("mcSister_3_mcPDG", "mcMother(mcDaughter(3, mcPDG))")
|
|
36
|
+
vm.addAlias("mcSister_4_mcPDG", "mcMother(mcDaughter(4, mcPDG))")
|
|
37
|
+
vm.addAlias("mcSister_5_mcPDG", "mcMother(mcDaughter(5, mcPDG))")
|
|
38
|
+
vm.addAlias("mcSister_6_mcPDG", "mcMother(mcDaughter(6, mcPDG))")
|
|
39
|
+
vm.addAlias("mcSister_7_mcPDG", "mcMother(mcDaughter(7, mcPDG))")
|
|
40
|
+
vm.addAlias("mcSister_8_mcPDG", "mcMother(mcDaughter(8, mcPDG))")
|
|
41
|
+
vm.addAlias("mcSister_9_mcPDG", "mcMother(mcDaughter(9, mcPDG))")
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def input_to_the_path():
|
|
45
|
+
ma.inputMdstList(
|
|
46
|
+
filelist=input_file_paths,
|
|
47
|
+
path=main,
|
|
48
|
+
environmentType="default",
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def reconstruct_generator_level():
|
|
53
|
+
ma.fillParticleListFromMC(decayString="K+:gen", cut="", path=main)
|
|
54
|
+
ma.fillParticleListFromMC(decayString="pi-:gen", cut="", path=main)
|
|
55
|
+
ma.fillParticleListFromMC(decayString="mu+:gen", cut="", path=main)
|
|
56
|
+
|
|
57
|
+
ma.reconstructMCDecay("K*0:gen =direct=> K+:gen pi-:gen", cut="", path=main)
|
|
58
|
+
ma.reconstructMCDecay("B0:gen =direct=> K*0:gen mu+:gen mu-:gen", cut="", path=main)
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def create_variable_lists():
|
|
62
|
+
std_vars = (
|
|
63
|
+
vc.deltae_mbc
|
|
64
|
+
+ vc.inv_mass
|
|
65
|
+
+ vc.mc_truth
|
|
66
|
+
+ ["mcMother_mcPDG"]
|
|
67
|
+
+ ["PDG"]
|
|
68
|
+
+ [
|
|
69
|
+
"mcSister_0_mcPDG",
|
|
70
|
+
"mcSister_1_mcPDG",
|
|
71
|
+
"mcSister_2_mcPDG",
|
|
72
|
+
"mcSister_3_mcPDG",
|
|
73
|
+
"mcSister_4_mcPDG",
|
|
74
|
+
"mcSister_5_mcPDG",
|
|
75
|
+
"mcSister_6_mcPDG",
|
|
76
|
+
"mcSister_7_mcPDG",
|
|
77
|
+
"mcSister_8_mcPDG",
|
|
78
|
+
"mcSister_9_mcPDG",
|
|
79
|
+
]
|
|
80
|
+
+ vc.pid
|
|
81
|
+
+ vc.kinematics
|
|
82
|
+
+ vc.mc_kinematics
|
|
83
|
+
+ ["dr", "dz"]
|
|
84
|
+
+ ["theta", "thetaErr", "mcTheta"]
|
|
85
|
+
+ ["isSignalAcceptBremsPhotons"]
|
|
86
|
+
)
|
|
87
|
+
|
|
88
|
+
Kstar0_vars = vu.create_aliases_for_selected(
|
|
89
|
+
list_of_variables=std_vars,
|
|
90
|
+
decay_string="B0 -> [^K*0 -> K+ pi-] mu+ mu-",
|
|
91
|
+
)
|
|
92
|
+
|
|
93
|
+
K_pi_vars = vu.create_aliases_for_selected(
|
|
94
|
+
list_of_variables=std_vars,
|
|
95
|
+
decay_string="B0 -> [K*0 -> ^K+ ^pi-] mu+ mu-",
|
|
96
|
+
prefix=["K_p", "pi_m"],
|
|
97
|
+
)
|
|
98
|
+
|
|
99
|
+
lepton_vars = vu.create_aliases_for_selected(
|
|
100
|
+
list_of_variables=std_vars,
|
|
101
|
+
decay_string="B0 -> [K*0 -> K+ pi-] ^mu+ ^mu-",
|
|
102
|
+
prefix=["mu_p", "mu_m"],
|
|
103
|
+
)
|
|
104
|
+
|
|
105
|
+
out = std_vars + Kstar0_vars + K_pi_vars + lepton_vars
|
|
106
|
+
return out
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def save_output(B0_vars, sim_level, file_path):
|
|
110
|
+
ma.variablesToNtuple(
|
|
111
|
+
decayString=f"B0:{sim_level}",
|
|
112
|
+
variables=B0_vars,
|
|
113
|
+
filename=file_path,
|
|
114
|
+
treename=sim_level,
|
|
115
|
+
path=main,
|
|
116
|
+
)
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
append_global_tag()
|
|
120
|
+
define_aliases()
|
|
121
|
+
input_to_the_path()
|
|
122
|
+
reconstruct_generator_level()
|
|
123
|
+
B0_vars = create_variable_lists()
|
|
124
|
+
save_output(B0_vars, "gen", output_file_path)
|
|
125
|
+
b2.process(main)
|
|
126
|
+
print(b2.statistics)
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
#!/usr/bin/env python3
|
|
2
|
+
# -*- coding: utf-8 -*-
|
|
3
|
+
|
|
4
|
+
# based off of Alexei's steering file
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
import sys
|
|
8
|
+
|
|
9
|
+
import basf2 as b2
|
|
10
|
+
import mdst as mdst
|
|
11
|
+
import glob as glob
|
|
12
|
+
|
|
13
|
+
dec_file_path = sys.argv[1]
|
|
14
|
+
out_file_path = sys.argv[2]
|
|
15
|
+
num_events = int(sys.argv[3])
|
|
16
|
+
|
|
17
|
+
print("\n-- Simulation Configuration --")
|
|
18
|
+
print(f"Decay file: {dec_file_path}")
|
|
19
|
+
print(f"Output file: {out_file_path}")
|
|
20
|
+
print(f"Number of events: {num_events}")
|
|
21
|
+
print("-----------------------\n")
|
|
22
|
+
|
|
23
|
+
# background (collision) files
|
|
24
|
+
bg = glob.glob(
|
|
25
|
+
"/group/belle2/dataprod/BGOverlay/early_phase3/release-06-00-05/overlay/BGx1/set0/*.root"
|
|
26
|
+
)
|
|
27
|
+
# background if running locally
|
|
28
|
+
bg_local = glob.glob(
|
|
29
|
+
"/group/belle2/dataprod/BGOverlay/early_phase3/release-06-00-05/overlay/BGx1/set0/*.root"
|
|
30
|
+
)
|
|
31
|
+
|
|
32
|
+
# set database conditions (in addition to default)
|
|
33
|
+
b2.conditions.prepend_globaltag("mc_production_MC15ri_a")
|
|
34
|
+
|
|
35
|
+
print("Create path")
|
|
36
|
+
main = b2.Path()
|
|
37
|
+
|
|
38
|
+
# default to early phase 3 (exp=1003), run 0
|
|
39
|
+
print("Add EventInfoSetter")
|
|
40
|
+
main.add_module("EventInfoSetter", expList=1003, runList=0, evtNumList=num_events)
|
|
41
|
+
|
|
42
|
+
# generate events from decfile
|
|
43
|
+
print("Add EvtGenInput")
|
|
44
|
+
main.add_module("EvtGenInput", userDECFile=b2.find_file(dec_file_path))
|
|
45
|
+
|
|
46
|
+
# Finally add mdst output (file name overwritten on the grid)
|
|
47
|
+
print("Add mdst output")
|
|
48
|
+
mdst.add_mdst_output(path=main, filename=str(out_file_path))
|
|
49
|
+
|
|
50
|
+
# process events and print call statistics
|
|
51
|
+
print("Process")
|
|
52
|
+
b2.process(path=main)
|
|
53
|
+
print(b2.statistics)
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
from helpers import (
|
|
2
|
+
Interval,
|
|
3
|
+
RunMetadata,
|
|
4
|
+
setup_run_dir,
|
|
5
|
+
submit_jobs,
|
|
6
|
+
)
|
|
7
|
+
|
|
8
|
+
# Setup
|
|
9
|
+
|
|
10
|
+
data_dir = "example/data/"
|
|
11
|
+
|
|
12
|
+
sim_steer_file_path = "example/steer_sim.py"
|
|
13
|
+
recon_steer_file_path = "example/steer_recon.py"
|
|
14
|
+
template_dec_file_path = "example/dec.dec"
|
|
15
|
+
|
|
16
|
+
parameter_bounds = {
|
|
17
|
+
"dC_7": Interval(0, 0),
|
|
18
|
+
"dC_9": Interval(-2, 1),
|
|
19
|
+
"dC_10": Interval(0, 0),
|
|
20
|
+
}
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
# Create a training dataset by sampling from a grid of parameters
|
|
24
|
+
|
|
25
|
+
training_run_metadata = RunMetadata(
|
|
26
|
+
split="train",
|
|
27
|
+
total_num_events=10_000,
|
|
28
|
+
num_trials=10, # each trial represents a different parameter configuration
|
|
29
|
+
num_subtrials_per_trial=1,
|
|
30
|
+
parameter_bounds=parameter_bounds,
|
|
31
|
+
sampling_type="grid",
|
|
32
|
+
parameter_grid_counts={ # number of grid points along each axis
|
|
33
|
+
"dC_7": 1,
|
|
34
|
+
"dC_9": 10,
|
|
35
|
+
"dC_10": 1,
|
|
36
|
+
},
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
training_run_dir = setup_run_dir(
|
|
40
|
+
run_metadata=training_run_metadata,
|
|
41
|
+
parent_dir=data_dir,
|
|
42
|
+
)
|
|
43
|
+
|
|
44
|
+
submit_jobs(
|
|
45
|
+
run_dir=training_run_dir,
|
|
46
|
+
sim_steer_file_path=sim_steer_file_path,
|
|
47
|
+
recon_steer_file_path=recon_steer_file_path,
|
|
48
|
+
template_dec_file_path=template_dec_file_path,
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
# Create a validation dataset by randomly sampling parameters
|
|
53
|
+
|
|
54
|
+
validation_run_metadata = RunMetadata(
|
|
55
|
+
split="val",
|
|
56
|
+
total_num_events=5_000,
|
|
57
|
+
num_trials=5,
|
|
58
|
+
num_subtrials_per_trial=2, # number of files per parameter value set
|
|
59
|
+
parameter_bounds=parameter_bounds,
|
|
60
|
+
sampling_type="random",
|
|
61
|
+
)
|
|
62
|
+
|
|
63
|
+
validation_run_dir = setup_run_dir(
|
|
64
|
+
run_metadata=validation_run_metadata, parent_dir=data_dir
|
|
65
|
+
)
|
|
66
|
+
|
|
67
|
+
submit_jobs(
|
|
68
|
+
run_dir=validation_run_dir,
|
|
69
|
+
sim_steer_file_path=sim_steer_file_path,
|
|
70
|
+
recon_steer_file_path=recon_steer_file_path,
|
|
71
|
+
template_dec_file_path=template_dec_file_path,
|
|
72
|
+
)
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling >= 1.26"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "b2_run_sim"
|
|
7
|
+
version = "0.0.1"
|
|
8
|
+
authors = [
|
|
9
|
+
{ name="Ethan Lee", email="elee20@hawaii.edu" },
|
|
10
|
+
]
|
|
11
|
+
description = "Simulation helper for basf2."
|
|
12
|
+
readme = "README.md"
|
|
13
|
+
requires-python = ">=3.14"
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"Operating System :: OS Independent",
|
|
17
|
+
]
|
|
18
|
+
license = "MIT"
|
|
19
|
+
license-files = ["LICEN[CS]E*"]
|
|
20
|
+
|
|
21
|
+
[project.urls]
|
|
22
|
+
Homepage = "https://github.com/ethanlee20/b2_sim"
|
|
23
|
+
Issues = "https://github.com/ethanlee20/b2_sim/issues"
|
|
@@ -0,0 +1,142 @@
|
|
|
1
|
+
from pathlib import Path
|
|
2
|
+
from dataclasses import dataclass, field
|
|
3
|
+
|
|
4
|
+
from .util import Interval, load_dataclass_from_json_file, save_dataclass_to_json_file
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
@dataclass
|
|
8
|
+
class Filenames:
|
|
9
|
+
metadata: Path | str = "metadata.json"
|
|
10
|
+
decay: Path | str = "decay.dec"
|
|
11
|
+
log: Path | str = "log.log"
|
|
12
|
+
|
|
13
|
+
def recon(self, subtrial: int):
|
|
14
|
+
out = f"recon_{subtrial}.root"
|
|
15
|
+
return out
|
|
16
|
+
|
|
17
|
+
def sim(self, subtrial: int):
|
|
18
|
+
out = f"sim_{subtrial}.root"
|
|
19
|
+
return out
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
@dataclass
|
|
23
|
+
class FilePaths:
|
|
24
|
+
dir_: Path | str
|
|
25
|
+
metadata: Path = field(init=False)
|
|
26
|
+
decay: Path = field(init=False)
|
|
27
|
+
log: Path = field(init=False)
|
|
28
|
+
|
|
29
|
+
def __post_init__(self):
|
|
30
|
+
self.dir_ = Path(self.dir_)
|
|
31
|
+
filenames = Filenames()
|
|
32
|
+
self.metadata = self.dir_.joinpath(filenames.metadata)
|
|
33
|
+
self.decay = self.dir_.joinpath(filenames.decay)
|
|
34
|
+
self.log = self.dir_.joinpath(filenames.log)
|
|
35
|
+
|
|
36
|
+
def recon(self, subtrial: int):
|
|
37
|
+
filename = Filenames().recon(subtrial)
|
|
38
|
+
out = self.dir_.joinpath(filename)
|
|
39
|
+
return out
|
|
40
|
+
|
|
41
|
+
def sim(self, subtrial: int):
|
|
42
|
+
filename = Filenames().sim(subtrial)
|
|
43
|
+
out = self.dir_.joinpath(filename)
|
|
44
|
+
return out
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
@dataclass
|
|
48
|
+
class SubtrialMetadata:
|
|
49
|
+
subtrial_num: int
|
|
50
|
+
num_events: int
|
|
51
|
+
parameter_values: dict[str, float]
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
@dataclass
|
|
55
|
+
class TrialMetadata:
|
|
56
|
+
trial_num: int
|
|
57
|
+
num_events: int
|
|
58
|
+
num_subtrials: int
|
|
59
|
+
parameter_values: dict[str, float]
|
|
60
|
+
|
|
61
|
+
@property
|
|
62
|
+
def num_events_per_subtrial(self) -> int:
|
|
63
|
+
out = safer_convert_to_int(self.num_events / self.num_subtrials)
|
|
64
|
+
return out
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
@dataclass
|
|
68
|
+
class RunMetadata:
|
|
69
|
+
split: str
|
|
70
|
+
total_num_events: int
|
|
71
|
+
num_trials: int
|
|
72
|
+
num_subtrials_per_trial: int
|
|
73
|
+
parameter_bounds: dict[str, Interval]
|
|
74
|
+
sampling_type: str
|
|
75
|
+
parameter_grid_counts: None | dict[str, int] = None
|
|
76
|
+
|
|
77
|
+
def __post_init__(self):
|
|
78
|
+
if self.sampling_type not in ("grid", "random"):
|
|
79
|
+
raise ValueError
|
|
80
|
+
if self.sampling_type == "grid" and self.parameter_grid_counts is None:
|
|
81
|
+
raise ValueError
|
|
82
|
+
if self.sampling_type != "grid" and self.parameter_grid_counts is not None:
|
|
83
|
+
raise ValueError
|
|
84
|
+
if (
|
|
85
|
+
self.parameter_grid_counts is not None
|
|
86
|
+
and self.total_num_grid_points != self.num_trials
|
|
87
|
+
):
|
|
88
|
+
raise ValueError
|
|
89
|
+
|
|
90
|
+
@property
|
|
91
|
+
def num_events_per_trial(self) -> int:
|
|
92
|
+
out = safer_convert_to_int(self.total_num_events / self.num_trials)
|
|
93
|
+
return out
|
|
94
|
+
|
|
95
|
+
@property
|
|
96
|
+
def num_events_per_subtrial(self) -> int:
|
|
97
|
+
out = safer_convert_to_int(
|
|
98
|
+
self.num_events_per_trial / self.num_subtrials_per_trial
|
|
99
|
+
)
|
|
100
|
+
return out
|
|
101
|
+
|
|
102
|
+
@property
|
|
103
|
+
def total_num_grid_points(self) -> None | int:
|
|
104
|
+
if self.parameter_grid_counts is None:
|
|
105
|
+
return None
|
|
106
|
+
out = prod(self.parameter_grid_counts.values())
|
|
107
|
+
return out
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
def make_trial_metadata_list(
|
|
111
|
+
run_metadata: RunMetadata, parameter_values_list: list[dict[str, float]]
|
|
112
|
+
) -> list[TrialMetadata]:
|
|
113
|
+
assert len(parameter_values_list) == run_metadata.num_trials
|
|
114
|
+
out = [
|
|
115
|
+
TrialMetadata(
|
|
116
|
+
trial_num,
|
|
117
|
+
run_metadata.num_events_per_trial,
|
|
118
|
+
run_metadata.num_subtrials_per_trial,
|
|
119
|
+
parameter_values,
|
|
120
|
+
)
|
|
121
|
+
for trial_num, parameter_values in enumerate(parameter_values_list)
|
|
122
|
+
]
|
|
123
|
+
return out
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
def save_metadata_to_dir(
|
|
127
|
+
metadata: TrialMetadata | RunMetadata, dir_path: Path | str
|
|
128
|
+
) -> None:
|
|
129
|
+
dir_path = Path(dir_path)
|
|
130
|
+
filename = Filenames().metadata
|
|
131
|
+
path = dir_path.joinpath(filename)
|
|
132
|
+
save_dataclass_to_json_file(metadata, path)
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
def load_metadata_from_dir(
|
|
136
|
+
metadata_cls, dir_path: Path | str
|
|
137
|
+
) -> RunMetadata | TrialMetadata:
|
|
138
|
+
dir_path = Path(dir_path)
|
|
139
|
+
filename = Filenames().metadata
|
|
140
|
+
path = dir_path.joinpath(filename)
|
|
141
|
+
out = load_dataclass_from_json_file(metadata_cls, path)
|
|
142
|
+
return out
|
|
@@ -0,0 +1,117 @@
|
|
|
1
|
+
from pathlib import Path
|
|
2
|
+
from random import seed, uniform
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
|
|
5
|
+
from .util import product_combine, Interval, linspace
|
|
6
|
+
from .config import (
|
|
7
|
+
TrialMetadata,
|
|
8
|
+
RunMetadata,
|
|
9
|
+
save_metadata_to_dir,
|
|
10
|
+
make_trial_metadata_list,
|
|
11
|
+
)
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def set_random_seed(seed_: int | None = None):
|
|
15
|
+
seed(seed_)
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
@dataclass
|
|
19
|
+
class RandomSampler:
|
|
20
|
+
parameter_bounds: dict[str, Interval]
|
|
21
|
+
|
|
22
|
+
def sample(
|
|
23
|
+
self,
|
|
24
|
+
n: int,
|
|
25
|
+
) -> list[dict[str, float]]:
|
|
26
|
+
out = [self._sample_once() for _ in range(n)]
|
|
27
|
+
return out
|
|
28
|
+
|
|
29
|
+
def _sample_once(
|
|
30
|
+
self,
|
|
31
|
+
) -> dict[str, float]:
|
|
32
|
+
out = {name: uniform(*bounds) for name, bounds in self.parameter_bounds.items()}
|
|
33
|
+
return out
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
@dataclass
|
|
37
|
+
class GridSampler:
|
|
38
|
+
parameter_bounds: dict[str, Interval]
|
|
39
|
+
|
|
40
|
+
def sample(self, parameter_counts: dict[str, int]) -> list[dict[str, float]]:
|
|
41
|
+
samples_per_wc = self._samples_per_wc(parameter_counts)
|
|
42
|
+
out = product_combine(samples_per_wc)
|
|
43
|
+
return out
|
|
44
|
+
|
|
45
|
+
def _samples_per_wc(
|
|
46
|
+
self,
|
|
47
|
+
parameter_counts: dict[str, int],
|
|
48
|
+
) -> dict[str, list[float]]:
|
|
49
|
+
assert parameter_counts.keys() == self.parameter_bounds.keys()
|
|
50
|
+
out = {
|
|
51
|
+
name: linspace(*self.parameter_bounds[name], count)
|
|
52
|
+
for name, count in parameter_counts.items()
|
|
53
|
+
}
|
|
54
|
+
return out
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def sample(run_metadata: RunMetadata):
|
|
58
|
+
sampler_cls = GridSampler if run_metadata.sampling_type == "grid" else RandomSampler
|
|
59
|
+
sampler = sampler_cls(run_metadata.parameter_bounds)
|
|
60
|
+
out = (
|
|
61
|
+
sampler.sample(run_metadata.parameter_grid_counts)
|
|
62
|
+
if run_metadata.sampling_type == "grid"
|
|
63
|
+
else sampler.sample(run_metadata.num_trials)
|
|
64
|
+
)
|
|
65
|
+
return out
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
def run_dir_name(split: str) -> str:
|
|
69
|
+
out = f"{split}"
|
|
70
|
+
return out
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def trial_dir_name(trial_num: int) -> str:
|
|
74
|
+
out = f"{trial_num}"
|
|
75
|
+
return out
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
def run_dir_path(split: str, parent_dir_path: Path | str) -> Path:
|
|
79
|
+
parent_dir_path = Path(parent_dir_path)
|
|
80
|
+
name = run_dir_name(split)
|
|
81
|
+
out = parent_dir_path.joinpath(name)
|
|
82
|
+
return out
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def trial_dir_path(trial_num: int, parent_dir_path: Path | str) -> Path:
|
|
86
|
+
parent_dir_path = Path(parent_dir_path)
|
|
87
|
+
name = trial_dir_name(trial_num)
|
|
88
|
+
out = parent_dir_path.joinpath(name)
|
|
89
|
+
return out
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def setup_trial_dir(trial_metadata: TrialMetadata, parent_dir_path: Path | str) -> None:
|
|
93
|
+
trial_num = trial_metadata.trial_num
|
|
94
|
+
dir_ = trial_dir_path(trial_num, parent_dir_path)
|
|
95
|
+
dir_.mkdir()
|
|
96
|
+
save_metadata_to_dir(
|
|
97
|
+
trial_metadata,
|
|
98
|
+
dir_,
|
|
99
|
+
)
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
def setup_run_dir(
|
|
103
|
+
run_metadata: RunMetadata,
|
|
104
|
+
parent_dir: Path | str,
|
|
105
|
+
) -> None:
|
|
106
|
+
split = run_metadata.split
|
|
107
|
+
dir_ = run_dir_path(split, parent_dir)
|
|
108
|
+
dir_.mkdir(parents=True)
|
|
109
|
+
save_metadata_to_dir(
|
|
110
|
+
run_metadata,
|
|
111
|
+
dir_,
|
|
112
|
+
)
|
|
113
|
+
parameter_values = sample(run_metadata)
|
|
114
|
+
trial_metadatas = make_trial_metadata_list(run_metadata, parameter_values)
|
|
115
|
+
for trial_metadata in trial_metadatas:
|
|
116
|
+
setup_trial_dir(trial_metadata, dir_)
|
|
117
|
+
return dir_
|
|
@@ -0,0 +1,202 @@
|
|
|
1
|
+
from pathlib import Path
|
|
2
|
+
from subprocess import run
|
|
3
|
+
from time import sleep
|
|
4
|
+
|
|
5
|
+
from .config import TrialMetadata, load_metadata_from_dir, FilePaths
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def read_dec_file(path: Path | str) -> str:
|
|
9
|
+
with open(path, "r", encoding="utf-8") as dec_file:
|
|
10
|
+
out = dec_file.read()
|
|
11
|
+
return out
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def format_template_dec_file_string(
|
|
15
|
+
str_: str, parameter_values: dict[str, float]
|
|
16
|
+
) -> str:
|
|
17
|
+
for parameter_name in parameter_values.keys():
|
|
18
|
+
if "{" f"{parameter_name}" "}" not in str_:
|
|
19
|
+
raise ValueError(
|
|
20
|
+
f"Parameter name not found in decay file string: {parameter_name}"
|
|
21
|
+
)
|
|
22
|
+
try:
|
|
23
|
+
out = str_.format(**parameter_values)
|
|
24
|
+
return out
|
|
25
|
+
except KeyError:
|
|
26
|
+
raise ValueError("Not all decay file parameters were specified?")
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def format_template_dec_file(path: Path | str, parameter_values: dict[str, float]):
|
|
30
|
+
text = read_dec_file(path)
|
|
31
|
+
out = format_template_dec_file_string(text, parameter_values)
|
|
32
|
+
return out
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def write_formatted_dec_file(
|
|
36
|
+
formatted_dec_file_path: Path | str,
|
|
37
|
+
template_dec_file_path: Path | str,
|
|
38
|
+
parameter_values: dict[str, float],
|
|
39
|
+
):
|
|
40
|
+
formatted_dec_file_string = format_template_dec_file(
|
|
41
|
+
template_dec_file_path, parameter_values
|
|
42
|
+
)
|
|
43
|
+
with open(formatted_dec_file_path, "w") as formatted_dec_file:
|
|
44
|
+
formatted_dec_file.write(formatted_dec_file_string)
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def make_sim_command_string(
|
|
48
|
+
num_events: int,
|
|
49
|
+
steer_file_path: Path | str,
|
|
50
|
+
decay_file_path: Path | str,
|
|
51
|
+
out_file_path: Path | str,
|
|
52
|
+
log_file_path: Path | str,
|
|
53
|
+
):
|
|
54
|
+
out = f"basf2 {steer_file_path} {decay_file_path} {out_file_path} {num_events} &>> {log_file_path}"
|
|
55
|
+
return out
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def make_recon_command_string(
|
|
59
|
+
steer_file_path: Path | str,
|
|
60
|
+
sim_file_path: Path | str,
|
|
61
|
+
out_file_path: Path | str,
|
|
62
|
+
log_file_path: Path | str,
|
|
63
|
+
):
|
|
64
|
+
out = f"basf2 {steer_file_path} {sim_file_path} {out_file_path} &>> {log_file_path}"
|
|
65
|
+
return out
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
def make_remove_sim_file_command_string(sim_file_path: Path | str):
|
|
69
|
+
out = f"rm {sim_file_path}"
|
|
70
|
+
return out
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def make_bsub_command(
|
|
74
|
+
num_events: int,
|
|
75
|
+
sim_steer_file_path: Path | str,
|
|
76
|
+
recon_steer_file_path: Path | str,
|
|
77
|
+
decay_file_path: Path | str,
|
|
78
|
+
sim_file_path: Path | str,
|
|
79
|
+
recon_file_path: Path | str,
|
|
80
|
+
log_file_path: Path | str,
|
|
81
|
+
queue: str,
|
|
82
|
+
):
|
|
83
|
+
sim_command = make_sim_command_string(
|
|
84
|
+
num_events=num_events,
|
|
85
|
+
steer_file_path=sim_steer_file_path,
|
|
86
|
+
decay_file_path=decay_file_path,
|
|
87
|
+
out_file_path=sim_file_path,
|
|
88
|
+
log_file_path=log_file_path,
|
|
89
|
+
)
|
|
90
|
+
recon_command = make_recon_command_string(
|
|
91
|
+
steer_file_path=recon_steer_file_path,
|
|
92
|
+
sim_file_path=sim_file_path,
|
|
93
|
+
out_file_path=recon_file_path,
|
|
94
|
+
log_file_path=log_file_path,
|
|
95
|
+
)
|
|
96
|
+
remove_sim_file_command = make_remove_sim_file_command_string(
|
|
97
|
+
sim_file_path=sim_file_path
|
|
98
|
+
)
|
|
99
|
+
|
|
100
|
+
out = f'bsub -q {queue} "{sim_command} && {recon_command} && {remove_sim_file_command}"'
|
|
101
|
+
return out
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
def submit_job(
|
|
105
|
+
num_events: int,
|
|
106
|
+
sim_steer_file_path: Path | str,
|
|
107
|
+
recon_steer_file_path: Path | str,
|
|
108
|
+
decay_file_path: Path | str,
|
|
109
|
+
sim_file_path: Path | str,
|
|
110
|
+
recon_file_path: Path | str,
|
|
111
|
+
log_file_path: Path | str,
|
|
112
|
+
queue: str,
|
|
113
|
+
debug: bool = False,
|
|
114
|
+
) -> None:
|
|
115
|
+
command = make_bsub_command(
|
|
116
|
+
num_events=num_events,
|
|
117
|
+
sim_steer_file_path=sim_steer_file_path,
|
|
118
|
+
recon_steer_file_path=recon_steer_file_path,
|
|
119
|
+
decay_file_path=decay_file_path,
|
|
120
|
+
sim_file_path=sim_file_path,
|
|
121
|
+
recon_file_path=recon_file_path,
|
|
122
|
+
log_file_path=log_file_path,
|
|
123
|
+
queue=queue,
|
|
124
|
+
)
|
|
125
|
+
if debug:
|
|
126
|
+
print(command, "\n")
|
|
127
|
+
return
|
|
128
|
+
run(command, shell=True)
|
|
129
|
+
|
|
130
|
+
|
|
131
|
+
def trial_dir_is_incomplete(dir_path: Path | str) -> bool:
|
|
132
|
+
try:
|
|
133
|
+
trial_metadata = load_metadata_from_dir(TrialMetadata, dir_path)
|
|
134
|
+
except FileNotFoundError:
|
|
135
|
+
raise RuntimeError(f"Trial directory does not contain metadata:\n{dir_path}")
|
|
136
|
+
num_subtrials = trial_metadata.num_subtrials
|
|
137
|
+
for subtrial in range(num_subtrials):
|
|
138
|
+
recon_file_path = FilePaths(dir_path).recon(subtrial)
|
|
139
|
+
if not recon_file_path.is_file():
|
|
140
|
+
return True
|
|
141
|
+
return False
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
def find_incomplete_trial_dirs(dir_path: Path | str) -> list[Path]:
|
|
145
|
+
dir_path = Path(dir_path)
|
|
146
|
+
out = [
|
|
147
|
+
child
|
|
148
|
+
for child in dir_path.iterdir()
|
|
149
|
+
if child.is_dir() and trial_dir_is_incomplete(child)
|
|
150
|
+
]
|
|
151
|
+
return out
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
def submit_jobs(
|
|
155
|
+
run_dir: Path | str,
|
|
156
|
+
sim_steer_file_path: Path | str,
|
|
157
|
+
recon_steer_file_path: Path | str,
|
|
158
|
+
template_dec_file_path: Path | str,
|
|
159
|
+
queue: str = "l",
|
|
160
|
+
batch_size: int = 200,
|
|
161
|
+
batch_wait_sec: int = 30,
|
|
162
|
+
job_wait_sec: int | float = 0.1,
|
|
163
|
+
verbose: bool = True,
|
|
164
|
+
debug: bool = False,
|
|
165
|
+
) -> None:
|
|
166
|
+
"""
|
|
167
|
+
Submit jobs.
|
|
168
|
+
|
|
169
|
+
Rerunning will skip trials deemed complete.
|
|
170
|
+
"""
|
|
171
|
+
|
|
172
|
+
submitted_job_count = 0
|
|
173
|
+
|
|
174
|
+
incomplete_trial_dirs = find_incomplete_trial_dirs(run_dir)
|
|
175
|
+
|
|
176
|
+
for trial_dir in incomplete_trial_dirs:
|
|
177
|
+
if verbose:
|
|
178
|
+
print(f"Submitting jobs for {trial_dir}.")
|
|
179
|
+
trial_metadata = load_metadata_from_dir(TrialMetadata, trial_dir)
|
|
180
|
+
trial_file_paths = FilePaths(trial_dir)
|
|
181
|
+
write_formatted_dec_file(
|
|
182
|
+
trial_file_paths.decay,
|
|
183
|
+
template_dec_file_path,
|
|
184
|
+
trial_metadata.parameter_values,
|
|
185
|
+
)
|
|
186
|
+
for subtrial in range(trial_metadata.num_subtrials):
|
|
187
|
+
submit_job(
|
|
188
|
+
num_events=trial_metadata.num_events_per_subtrial,
|
|
189
|
+
sim_steer_file_path=sim_steer_file_path,
|
|
190
|
+
recon_steer_file_path=recon_steer_file_path,
|
|
191
|
+
decay_file_path=trial_file_paths.decay,
|
|
192
|
+
sim_file_path=trial_file_paths.sim(subtrial),
|
|
193
|
+
recon_file_path=trial_file_paths.recon(subtrial),
|
|
194
|
+
log_file_path=trial_file_paths.log,
|
|
195
|
+
queue=queue,
|
|
196
|
+
debug=debug,
|
|
197
|
+
)
|
|
198
|
+
submitted_job_count += 1
|
|
199
|
+
sleep(job_wait_sec)
|
|
200
|
+
|
|
201
|
+
if submitted_job_count % batch_size == 0:
|
|
202
|
+
sleep(batch_wait_sec)
|
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
from typing import Any
|
|
2
|
+
from itertools import product
|
|
3
|
+
from math import prod
|
|
4
|
+
from dataclasses import dataclass, asdict, field, fields, is_dataclass
|
|
5
|
+
from copy import deepcopy
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
from json import load, dump
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def product_combine(d: dict[Any, list | tuple]):
|
|
11
|
+
prod = product(*d.values())
|
|
12
|
+
out = [dict(zip(d.keys(), i)) for i in prod]
|
|
13
|
+
return out
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def safer_convert_to_int(
|
|
17
|
+
x: float,
|
|
18
|
+
) -> int:
|
|
19
|
+
assert x.is_integer()
|
|
20
|
+
return int(x)
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def load_json(
|
|
24
|
+
path: Path | str,
|
|
25
|
+
) -> dict:
|
|
26
|
+
|
|
27
|
+
path = Path(path)
|
|
28
|
+
if not path.is_file():
|
|
29
|
+
raise FileNotFoundError
|
|
30
|
+
with open(path, "r") as f:
|
|
31
|
+
out = load(f)
|
|
32
|
+
return out
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def dump_json(
|
|
36
|
+
obj: dict,
|
|
37
|
+
path: Path | str,
|
|
38
|
+
) -> None:
|
|
39
|
+
path = Path(path)
|
|
40
|
+
with open(path, "x") as f:
|
|
41
|
+
return dump(
|
|
42
|
+
obj,
|
|
43
|
+
f,
|
|
44
|
+
indent=4,
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def read_from_nested_dict(dict_, keys: list):
|
|
49
|
+
for key in keys:
|
|
50
|
+
dict_ = dict_[key]
|
|
51
|
+
return dict_
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def write_to_nested_dict(dict_, keys: list, value):
|
|
55
|
+
dict_ = read_from_nested_dict(dict_, keys[:-1])
|
|
56
|
+
dict_[keys[-1]] = value
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def _rebuild_dataclass(cls, dict_: dict, keys: list | None = None):
|
|
60
|
+
if not is_dataclass(cls):
|
|
61
|
+
return
|
|
62
|
+
if keys is None:
|
|
63
|
+
keys = []
|
|
64
|
+
for field_ in fields(cls):
|
|
65
|
+
_rebuild_dataclass(field_.type, dict_, keys=keys + [field_.name])
|
|
66
|
+
if keys == []:
|
|
67
|
+
out = cls(**dict_)
|
|
68
|
+
return out
|
|
69
|
+
write_to_nested_dict(dict_, keys, cls(**read_from_nested_dict(dict_, keys)))
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def rebuild_dataclass(cls, dict_: dict):
|
|
73
|
+
"""
|
|
74
|
+
dict_ must not contain dataclasses.
|
|
75
|
+
"""
|
|
76
|
+
dict_ = deepcopy(dict_)
|
|
77
|
+
out = _rebuild_dataclass(cls, dict_)
|
|
78
|
+
return out
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def load_dataclass_from_json_file(cls: Any, path: Path | str):
|
|
82
|
+
dict_ = load_json(path)
|
|
83
|
+
out = rebuild_dataclass(cls, dict_)
|
|
84
|
+
return out
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def save_dataclass_to_json_file(obj, path):
|
|
88
|
+
dict_ = asdict(obj)
|
|
89
|
+
dump_json(dict_, path)
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def linspace(start: float, stop: float, num_points: int) -> list[float, ...]:
|
|
93
|
+
if num_points < 1:
|
|
94
|
+
raise ValueError("Need at least two points." f" Got: {num_points} points")
|
|
95
|
+
if num_points == 1:
|
|
96
|
+
if start != stop:
|
|
97
|
+
raise ValueError("If creating only one point, start must equal stop.")
|
|
98
|
+
out = [
|
|
99
|
+
start,
|
|
100
|
+
]
|
|
101
|
+
return out
|
|
102
|
+
num_spaces = num_points - 1
|
|
103
|
+
spacing = (stop - start) / num_spaces
|
|
104
|
+
out = [start + i * spacing for i in range(num_points)]
|
|
105
|
+
assert len(out) == num_points
|
|
106
|
+
assert out[-1] == stop
|
|
107
|
+
return out
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
@dataclass
|
|
111
|
+
class Interval:
|
|
112
|
+
left: float
|
|
113
|
+
right: float
|
|
114
|
+
|
|
115
|
+
def as_tuple(self):
|
|
116
|
+
return (self.left, self.right)
|
|
117
|
+
|
|
118
|
+
def __iter__(self):
|
|
119
|
+
tuple_ = self.as_tuple()
|
|
120
|
+
out = tuple_.__iter__()
|
|
121
|
+
return out
|