153 lines
4.6 KiB
Python
Executable File
153 lines
4.6 KiB
Python
Executable File
#!/usr/bin/env python
|
|
|
|
|
|
import jinja2
|
|
import json
|
|
import sys
|
|
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Tuple
|
|
|
|
|
|
SIZE = (192, 192, 192)
|
|
|
|
|
|
templates_env = jinja2.Environment(
|
|
loader=jinja2.FileSystemLoader(Path(__file__).parent.parent / "templates"),
|
|
autoescape=jinja2.select_autoescape()
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class Experiment:
|
|
job_name: str
|
|
account: str
|
|
partition: str
|
|
nastja_binary_path: str
|
|
nodes: int
|
|
tasks: int
|
|
num_blocks: Tuple[int, int, int]
|
|
domain_scale: Tuple[int, int, int]
|
|
time: str = "00:15:00"
|
|
extra_sbatch_line: str = ""
|
|
logfile_path: str = "/p/project/cellsinsilico/paulslustigebude/ma/experiments/eval/logs/%x-%A.%a"
|
|
config_path: str = "/p/project/cellsinsilico/paulslustigebude/ma/experiments/eval/generated/config/${SLURM_JOB_NAME}.json"
|
|
output_dir_path: str = "/p/scratch/cellsinsilico/paul/nastja-out/${SLURM_JOB_NAME}-${SLURM_ARRAY_JOB_ID}.${SLURM_ARRAY_TASK_ID}"
|
|
|
|
def get_config(self):
|
|
with (Path(__file__).parent.parent / "templates" / "weak.json").open(encoding="utf8") as f:
|
|
config = json.load(f)
|
|
|
|
size = (
|
|
SIZE[0] * self.domain_scale[0],
|
|
SIZE[1] * self.domain_scale[1],
|
|
SIZE[2] * self.domain_scale[2],
|
|
)
|
|
blocksize = (
|
|
size[0] // self.num_blocks[0],
|
|
size[1] // self.num_blocks[1],
|
|
size[2] // self.num_blocks[2],
|
|
)
|
|
config["Geometry"] = {
|
|
"blockcount": list(self.num_blocks),
|
|
"blocksize": list(blocksize),
|
|
}
|
|
|
|
cells_filling = [{
|
|
"box": [
|
|
[0, 0, 0],
|
|
list(size)
|
|
],
|
|
"celltype": 0,
|
|
"component": 0,
|
|
"pattern": "const",
|
|
"seed": 0,
|
|
"shape": "cube",
|
|
"value": 0,
|
|
}]
|
|
for z in range(self.domain_scale[2]):
|
|
for y in range(self.domain_scale[1]):
|
|
for x in range(self.domain_scale[0]):
|
|
cx = x * SIZE[0] + SIZE[0] // 2
|
|
cy = y * SIZE[1] + SIZE[1] // 2
|
|
cz = z * SIZE[2] + SIZE[2] // 2
|
|
cells_filling.append({
|
|
"shape": "sphere",
|
|
"pattern": "voronoi",
|
|
"count": 715,
|
|
"radius": 38,
|
|
"center": [cx, cy, cz],
|
|
"box": [
|
|
[cx - 38, cy - 38, cz - 38],
|
|
[cx + 38, cy + 38, cz + 38]
|
|
],
|
|
"celltype": 9,
|
|
"seed": 758960,
|
|
})
|
|
|
|
config["Filling"]["cells"] = cells_filling
|
|
|
|
return config
|
|
|
|
def write_batch_file(self, out_path: Path):
|
|
t = templates_env.get_template("strong-batch.j2")
|
|
t.stream(
|
|
name=self.job_name,
|
|
account=self.account,
|
|
partition=self.partition,
|
|
nodes=self.nodes,
|
|
tasks=self.tasks,
|
|
extra_sbatch_line=self.extra_sbatch_line,
|
|
time=self.time,
|
|
logfile_path=self.logfile_path,
|
|
nastja_binary_path=self.nastja_binary_path,
|
|
config_path=self.config_path,
|
|
output_dir_path=self.output_dir_path,
|
|
).dump(str(out_path))
|
|
|
|
|
|
def make_cpu_ex(x: int, y: int, z: int) -> Experiment:
|
|
num_blocks = x * y * z
|
|
assert num_blocks % 48 == 0
|
|
num_nodes = num_blocks // 48
|
|
|
|
assert x % 4 == 0
|
|
assert y % 4 == 0
|
|
assert z % 3 == 0
|
|
|
|
return Experiment(
|
|
job_name=f"weak-cpu-{x:02}-{y:02}-{z:02}",
|
|
account="cellsinsilico",
|
|
partition="batch",
|
|
nastja_binary_path="/p/project/cellsinsilico/paulslustigebude/nastja/build-nocuda/nastja",
|
|
nodes=num_nodes,
|
|
tasks=num_blocks,
|
|
num_blocks=(x, y, z),
|
|
domain_scale=(x // 4, y // 4, z // 3),
|
|
)
|
|
|
|
|
|
experiments = [
|
|
make_cpu_ex(4, 4, 3),
|
|
make_cpu_ex(4, 4, 6),
|
|
make_cpu_ex(4, 4, 12),
|
|
make_cpu_ex(4, 8, 12),
|
|
make_cpu_ex(8, 8, 12),
|
|
make_cpu_ex(8, 8, 24),
|
|
make_cpu_ex(8, 16, 24),
|
|
make_cpu_ex(16, 16, 24),
|
|
]
|
|
|
|
if __name__ == "__main__":
|
|
outdir = Path(__file__).parent.parent / "generated"
|
|
|
|
for e in experiments:
|
|
print(f"Generating config for {e.job_name}", file=sys.stderr)
|
|
config_path = (outdir / "config" / e.job_name).with_suffix(".json")
|
|
with config_path.open("w", encoding="utf8") as f:
|
|
json.dump(e.get_config(), f, indent=2)
|
|
print(f"Generating batch file for {e.job_name}", file=sys.stderr)
|
|
e.write_batch_file(outdir / "batch" / e.job_name)
|