File size: 2,964 Bytes
e2703dc
 
ede74c0
 
e2703dc
 
ede74c0
e2703dc
 
 
ede74c0
e2703dc
 
 
 
 
 
ede74c0
 
e2703dc
ede74c0
 
 
e2703dc
 
ede74c0
 
 
e2703dc
 
ede74c0
 
 
 
 
 
e2703dc
 
ede74c0
 
 
 
e2703dc
 
ede74c0
 
e2703dc
 
ede74c0
 
 
 
e2703dc
 
ede74c0
 
e2703dc
ede74c0
 
 
 
 
 
 
 
e2703dc
 
ede74c0
 
 
 
 
 
 
e2703dc
 
ede74c0
 
 
 
 
 
 
 
 
 
 
 
e2703dc
 
ede74c0
 
 
 
 
e2703dc
 
ede74c0
 
 
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
from __future__ import annotations

import json
import random
import sys
from pathlib import Path
from typing import Any

import numpy as np
import torch
import yaml

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


def load_config(config_path: str | Path | None = None) -> dict[str, Any]:
    path = Path(config_path) if config_path else PROJECT_ROOT / "conf" / "config.yaml"
    if not path.is_absolute():
        path = PROJECT_ROOT / path
    with path.open("r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def project_path(path: str | Path) -> Path:
    path = Path(path)
    return path if path.is_absolute() else PROJECT_ROOT / path


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def select_device(value: str = "auto") -> torch.device:
    if value == "auto":
        return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    return torch.device(value)


def ensure_parent(path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)


def ensure_runtime_dirs(cfg: dict[str, Any]) -> None:
    for key in ("fake_data", "prediction", "metrics", "loss", "figure"):
        ensure_parent(project_path(cfg["paths"][key]))
    ensure_parent(project_path(cfg["training"]["checkpoint"]))


def build_model(cfg: dict[str, Any]) -> torch.nn.Module:
    from model import PINNsformer1D

    model_cfg = cfg["model"]
    return PINNsformer1D(
        d_out=int(model_cfg["d_out"]),
        d_hidden=int(model_cfg["d_hidden"]),
        d_model=int(model_cfg["d_model"]),
        N=int(model_cfg["num_layers"]),
        heads=int(model_cfg["heads"]),
    )


def initial_condition(x: np.ndarray | torch.Tensor, cfg: dict[str, Any]):
    initial = cfg["equation"]["initial"]
    center = initial["center"]
    sigma = initial["sigma"]
    if isinstance(x, torch.Tensor):
        return torch.exp(-((x - center) ** 2) / (2 * sigma**2))
    return np.exp(-((x - center) ** 2) / (2 * sigma**2))


def exact_reaction_solution(
    x: np.ndarray | torch.Tensor,
    t: np.ndarray | torch.Tensor,
    cfg: dict[str, Any],
):
    h = initial_condition(x, cfg)
    rate = cfg["equation"]["reaction_rate"]
    if isinstance(h, torch.Tensor):
        exp_term = torch.exp(rate * t)
    else:
        exp_term = np.exp(rate * t)
    return h * exp_term / (h * exp_term + 1 - h)


def relative_errors(pred: np.ndarray, target: np.ndarray) -> dict[str, float]:
    return {
        "relative_l1": float(np.sum(np.abs(target - pred)) / np.sum(np.abs(target))),
        "relative_l2": float(np.sqrt(np.sum((target - pred) ** 2) / np.sum(target**2))),
    }


def write_json(path: Path, payload: dict[str, Any]) -> None:
    ensure_parent(path)
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")