Domain Mixture Scaling Law

from __future__ import annotations
from typing import List, Dict
import math

# Log-scaling law per group and domain:
# loss_domain_i = a + b * log(proportion_domain_i + eps)
# The functional form is identical across groups; only (a, b, eps) differ.

COEFFS: Dict[str, Dict[str, Dict[str, float]]] = {
    "160M": {
        "loss_domain_1": {"a": 2.2532615943868906, "b": -0.1337057136545849, "eps": 0.0023901343745956435},
        "loss_domain_2": {"a": 3.2631919363951556, "b": -0.03558508713341269, "eps": 0.0028180383737978713},
        "loss_domain_3": {"a": 2.617689579679002, "b": -0.09252884608244097, "eps": 0.000733714154677849},
        "loss_domain_4": {"a": 1.1948582730723922, "b": -0.12312317609847837, "eps": 0.00195132901468823},
        "loss_domain_5": {"a": 3.0846493236033896, "b": -0.14154552771058287, "eps": 0.027307052357303876},
    },
    "305M": {
        "loss_domain_1": {"a": 2.1166886286314748, "b": -0.12677944713435335, "eps": 0.002106041895388586},
        "loss_domain_2": {"a": 3.0993803822936004, "b": -0.0385714100210861, "eps": 0.00469352643305276},
        "loss_domain_3": {"a": 2.498051621939985, "b": -0.08371626119596315, "eps": 0.0003880182070929469},
        "loss_domain_4": {"a": 1.0998534285937238, "b": -0.11664800642521109, "eps": 0.001864178447017992},
        "loss_domain_5": {"a": 2.918795681535315, "b": -0.13713177186385297, "eps": 0.023253743472940992},
    },
    "410M": {
        "loss_domain_1": {"a": 2.0637659220766347, "b": -0.12150735957421108, "eps": 0.0017973751820540681},
        "loss_domain_2": {"a": 3.0488279748294658, "b": -0.027049899121057236, "eps": 0.0012212145195230296},
        "loss_domain_3": {"a": 2.431248081266068, "b": -0.08424573100842837, "eps": 0.00036447069913629604},
        "loss_domain_4": {"a": 1.0617583927188603, "b": -0.11564152762782738, "eps": 0.0020181562979846223},
        "loss_domain_5": {"a": 2.8417472866298, "b": -0.13918681081923803, "eps": 0.02183585057099095},
    },
    "70M": {
        "loss_domain_1": {"a": 2.536901150373465, "b": -0.1533997450030275, "eps": 0.0032688441418023727},
        "loss_domain_2": {"a": 3.568955616082981, "b": -0.05120827266924541, "eps": 0.007660059062331851},
        "loss_domain_3": {"a": 2.886179438825381, "b": -0.10197118388034185, "eps": 0.0009064844324661106},
        "loss_domain_4": {"a": 1.4039456278566775, "b": -0.13849029661959558, "eps": 0.0019748899726039115},
        "loss_domain_5": {"a": 3.435925433424547, "b": -0.1375046497033619, "eps": 0.02690444281470573},
    },
}

_PROP_COLS = [
    "proportion_domain_1",
    "proportion_domain_2",
    "proportion_domain_3",
    "proportion_domain_4",
    "proportion_domain_5",
]
_LOSS_COLS = [
    "loss_domain_1",
    "loss_domain_2",
    "loss_domain_3",
    "loss_domain_4",
    "loss_domain_5",
]


def _predict_for_group(row: Dict[str, float], group: str) -> Dict[str, float]:
    coeffs = COEFFS.get(group)
    if coeffs is None:
        raise ValueError(f"Unknown group '{group}'. Known groups: {sorted(COEFFS)}")
    out: Dict[str, float] = {}
    for i, loss_key in enumerate(_LOSS_COLS):
        prop_key = _PROP_COLS[i]
        p = float(row.get(prop_key, 0.0))
        a = coeffs[loss_key]["a"]
        b = coeffs[loss_key]["b"]
        eps = coeffs[loss_key]["eps"]
        # Guard against tiny negatives from numerical issues
        val = a + b * math.log(max(p + eps, 1e-12))
        out[loss_key] = float(val)
    return out


def law(input_data: List[Dict[str, float]], group: str) -> List[Dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    return [_predict_for_group(row, group) for row in input_data]

from __future__ import annotations
from typing import Dict, List

# Discovered scaling law (same functional form across groups):
#   loss_domain_i = a_{g,i} + b_{g,i} * (proportion_domain_i + c_{g,i}) ** (-alpha_{g,i})
# Parameters (a, b, alpha, c) are fitted per group g and domain i.

_PARAMS: Dict[str, Dict[int, Dict[str, float]]] = {
    "160M": {
        1: {"a": -0.11147193199439992, "b": 2.3787691982958217, "alpha": 0.05, "c": 0.003125},
        2: {"a": 3.278500369142649,   "b": 0.01306711991811282, "alpha": 0.8034482758620689, "c": 0.034895833333333334},
        3: {"a": 0.8506733951924391,  "b": 1.7622754024333271, "alpha": 0.05, "c": 0.0015625},
        4: {"a": -1.072627406051137,  "b": 2.2748208460279304, "alpha": 0.05, "c": 0.003125},
        5: {"a": 0.3144198553952366,  "b": 2.772096693493694,  "alpha": 0.05, "c": 0.034895833333333334},
    },
    "305M": {
        1: {"a": -0.18733739807046243, "b": 2.312924189576437,  "alpha": 0.05, "c": 0.003125},
        2: {"a": 3.0583843406593463,  "b": 0.05867924234335276, "alpha": 0.3672413793103448, "c": 0.019791666666666666},
        3: {"a": 2.037098981893403,   "b": 0.48486293207181885, "alpha": 0.1293103448275862, "c": 0.0015625},
        4: {"a": -1.0727458525278075, "b": 2.177270392291824,  "alpha": 0.05, "c": 0.003125},
        5: {"a": 1.465292570110724,   "b": 1.4569415752353443, "alpha": 0.0896551724137931, "c": 0.034895833333333334},
    },
    "410M": {
        1: {"a": -0.21860503605523499, "b": 2.2855798194974324, "alpha": 0.05, "c": 0.003125},
        2: {"a": 3.068706883350695,    "b": 0.005605907282164662, "alpha": 1.0017241379310344, "c": 0.034895833333333334},
        3: {"a": 1.958790095134566,    "b": 0.4939182316775028,  "alpha": 0.1293103448275862, "c": 0.0015625},
        4: {"a": -1.050236886559304,   "b": 2.1205339628269058, "alpha": 0.05, "c": 0.003125},
        5: {"a": 1.3466574766335897,   "b": 1.4981064795012045,  "alpha": 0.0896551724137931, "c": 0.034895833333333334},
    },
    "70M": {
        1: {"a": 2.1270757115611185,  "b": 0.4395199897525094,  "alpha": 0.2482758620689655, "c": 0.013194444444444443},
        2: {"a": 3.552025532207962,   "b": 0.04562819235160885, "alpha": 0.5258620689655172, "c": 0.034895833333333334},
        3: {"a": 1.0433037320725291,  "b": 1.8510839638322976,  "alpha": 0.05, "c": 0.0015625},
        4: {"a": 1.3203779743559676,  "b": 0.16265793539818846, "alpha": 0.4068965517241379, "c": 0.013194444444444443},
        5: {"a": 3.192722657529413,   "b": 0.2586205571225181,  "alpha": 0.446551724137931, "c": 0.09531249999999998},
    },
}

_DOMAIN_KEYS = {
    1: ("proportion_domain_1", "loss_domain_1"),
    2: ("proportion_domain_2", "loss_domain_2"),
    3: ("proportion_domain_3", "loss_domain_3"),
    4: ("proportion_domain_4", "loss_domain_4"),
    5: ("proportion_domain_5", "loss_domain_5"),
}

_DEF_EPS = 1e-8  # numerical floor to avoid zero to negative-power


def _predict_for_group_row(row: Dict[str, float], params_g: Dict[int, Dict[str, float]]) -> Dict[str, float]:
    out: Dict[str, float] = {}
    for i in range(1, 6):
        p_key, y_key = _DOMAIN_KEYS[i]
        p = float(row.get(p_key, 0.0))
        par = params_g[i]
        # Apply shifted power law with small numerical floor
        val = par["a"] + par["b"] * ((max(p, 0.0) + max(par["c"], _DEF_EPS)) ** (-par["alpha"]))
        out[y_key] = float(val)
    return out


def law(input_data: List[Dict[str, float]], group: str) -> List[Dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    if group not in _PARAMS:
        raise ValueError(f"Unknown group '{group}'. Known: {sorted(_PARAMS.keys())}")
    params_g = _PARAMS[group]
    return [_predict_for_group_row(row, params_g) for row in input_data]

from typing import List, Dict
import math

# Offset power-law with asymptote per group and domain:
# loss_domain_i = a + b * (proportion_domain_i + d) ** c
# Parameters were fitted on the provided dataset.
_PARAMS = {
    "70M": {
        1: {"a": 2.1328, "b": 0.4344, "c": -0.25, "d": 0.01320},
        2: {"a": 3.5859, "b": 0.0212, "c": -0.80, "d": 0.05000},
        3: {"a": 2.6248, "b": 0.3145, "c": -0.20, "d": 0.00349},
        4: {"a": 1.3108, "b": 0.1693, "c": -0.40, "d": 0.01320},
        5: {"a": 2.8438, "b": 0.6006, "c": -0.20, "d": 0.05000},
    },
    "160M": {
        1: {"a": 1.7899, "b": 0.4896, "c": -0.20, "d": 0.00847},
        2: {"a": 3.2759, "b": 0.0149, "c": -0.75, "d": 0.03208},
        3: {"a": 2.4432, "b": 0.2486, "c": -0.20, "d": 0.00224},
        4: {"a": 0.9249, "b": 0.3149, "c": -0.25, "d": 0.00847},
        5: {"a": 2.4777, "b": 0.6157, "c": -0.20, "d": 0.05000},
    },
    "305M": {
        1: {"a": 1.6614, "b": 0.4760, "c": -0.20, "d": 0.00847},
        2: {"a": 3.1208, "b": 0.0108, "c": -0.95, "d": 0.05000},
        3: {"a": 2.3548, "b": 0.2163, "c": -0.20, "d": 0.00143},
        4: {"a": 0.8390, "b": 0.3014, "c": -0.25, "d": 0.00847},
        5: {"a": 2.3085, "b": 0.6178, "c": -0.20, "d": 0.05000},
    },
    "410M": {
        1: {"a": 1.6083, "b": 0.4704, "c": -0.20, "d": 0.00847},
        2: {"a": 3.0372, "b": 0.0275, "c": -0.45, "d": 0.01320},
        3: {"a": 2.2824, "b": 0.2203, "c": -0.20, "d": 0.00143},
        4: {"a": 0.8117, "b": 0.2936, "c": -0.25, "d": 0.00847},
        5: {"a": 2.2138, "b": 0.6352, "c": -0.20, "d": 0.05000},
    },
}

# Fallback parameters: average across known groups per domain.
# This is used if an unknown group is requested.
_avg_params = {}
for dom in range(1, 6):
    vals = [g[dom] for g in _PARAMS.values()]
    _avg_params[dom] = {
        "a": sum(v["a"] for v in vals) / len(vals),
        "b": sum(v["b"] for v in vals) / len(vals),
        "c": sum(v["c"] for v in vals) / len(vals),
        "d": sum(v["d"] for v in vals) / len(vals),
    }


def _predict_for_group(row: Dict[str, float], group: str) -> Dict[str, float]:
    params = _PARAMS.get(group, _avg_params)
    out: Dict[str, float] = {}
    for i in range(1, 6):
        p = float(row.get(f"proportion_domain_{i}", 0.0))
        if p < 0.0:
            p = 0.0
        pr = params[i]
        a = pr["a"]
        b = pr["b"]
        c = pr["c"]
        d = pr["d"]
        y = a + b * ((p + d) ** c)
        out[f"loss_domain_{i}"] = float(y)
    return out


def law(input_data: List[Dict[str, float]], group: str) -> List[Dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    return [_predict_for_group(row, group) for row in input_data]

from __future__ import annotations

import math
from typing import Dict, List

# Small epsilon to avoid log(0) when a domain proportion is zero
_EPS = 1e-6

# Per-group, per-output coefficients for the multi-log linear law.
# For each group and each output loss_domain_k, the list is:
#   [intercept, b1, b2, b3, b4, b5]
_COEFS: Dict[str, Dict[str, List[float]]] = {
    "160M": {
        "loss_domain_1": [
            2.4042025803406553,
            -0.04878456902733,
            -0.001374141990715709,
            0.00425143168670042,
            -0.004036425589256114,
            0.007660516801797778,
        ],
        "loss_domain_2": [
            3.329935868696654,
            0.002613273350161564,
            -0.011891996557419415,
            2.5572674116154508e-05,
            0.003256629390122456,
            -0.005952205705195364,
        ],
        "loss_domain_3": [
            2.8006343825165314,
            0.0038463371713819104,
            0.0010271048724917586,
            -0.03834514983761013,
            0.003370871701657902,
            -0.004992444549614571,
        ],
        "loss_domain_4": [
            1.3842254371292473,
            0.0001770125184695076,
            0.0019008052563317707,
            0.00040985897229494966,
            -0.04321725706040489,
            0.005203775943861618,
        ],
        "loss_domain_5": [
            3.22070848376419,
            0.005681276339156933,
            0.004826630941776033,
            -0.00156340623248917,
            0.005473022220321205,
            -0.030039287316337895,
        ],
    },
    "305M": {
        "loss_domain_1": [
            2.254097761686427,
            -0.047618308465165196,
            -0.0013550046237553216,
            0.004003192579882288,
            -0.0042909223510698554,
            0.006758159630738282,
        ],
        "loss_domain_2": [
            3.166848403701224,
            0.0023474329227610603,
            -0.011422287483401284,
            4.890168638858063e-06,
            0.0026987505976661923,
            -0.006498741023938888,
        ],
        "loss_domain_3": [
            2.6500359123770685,
            0.0035548451771600445,
            0.0011888369505355075,
            -0.03964115398330717,
            0.0028959099742507494,
            -0.004828358464409514,
        ],
        "loss_domain_4": [
            1.2751944568968825,
            -0.00024923215104993846,
            0.0015693015233153626,
            0.0004352251902864513,
            -0.041431801189773,
            0.00487564464080058,
        ],
        "loss_domain_5": [
            3.0477226649403506,
            0.005424541586794049,
            0.004752807485933067,
            -0.0015473014680650466,
            0.004775541162988269,
            -0.030879417992197157,
        ],
    },
    "410M": {
        "loss_domain_1": [
            2.1969109104720936,
            -0.04711053566759858,
            -0.0005929598868873091,
            0.0032935545744610275,
            -0.0034299883497122844,
            0.006245418652814477,
        ],
        "loss_domain_2": [
            3.096774846113359,
            0.002683578481767577,
            -0.01079532012122322,
            -0.0006061709329291659,
            0.003289915492216653,
            -0.007059325448456723,
        ],
        "loss_domain_3": [
            2.5854698484676923,
            0.0036493802056689617,
            0.0014981486567392278,
            -0.04057579585320029,
            0.003422309711002145,
            -0.005304641068008599,
        ],
        "loss_domain_4": [
            1.2374283294627488,
            -0.0007019754734980677,
            0.0023220834389344196,
            0.00022483249867104974,
            -0.040255057354183485,
            0.0043998574943246395,
        ],
        "loss_domain_5": [
            2.9734625677861892,
            0.005474979456664107,
            0.005292433569369652,
            -0.0021983339298417347,
            0.005328130514437292,
            -0.031848404298810445,
        ],
    },
    "70M": {
        "loss_domain_1": [
            2.709753097020149,
            -0.05223448519596348,
            -0.002096832291136182,
            0.0049309629806633245,
            -0.005121637241168704,
            0.008582633603168494,
        ],
        "loss_domain_2": [
            3.6589940901202946,
            0.0024594854221674217,
            -0.013037306540415955,
            0.00018013301869309603,
            0.0025694674548725075,
            -0.0055192991657138785,
        ],
        "loss_domain_3": [
            3.0882519411499176,
            0.004138978755961604,
            0.000999612740187218,
            -0.040138664176247195,
            0.002610113875372775,
            -0.00489486237526904,
        ],
        "loss_domain_4": [
            1.6123830068952152,
            -0.0002626923872205353,
            0.001840626589019361,
            0.0004871656033250259,
            -0.048580221686945084,
            0.0053741041526707105,
        ],
        "loss_domain_5": [
            3.567736623439916,
            0.005906206376564464,
            0.00463119283202685,
            -0.0015851069219014631,
            0.004999662780701288,
            -0.029220774259506147,
        ],
    },
}


def _predict_one(input_point: Dict[str, float], beta: List[float]) -> float:
    # Build feature vector: [1, log(p1+eps), ..., log(p5+eps)]
    logs = [math.log(max(input_point.get(f"proportion_domain_{i}", 0.0), 0.0) + _EPS) for i in range(1, 6)]
    return beta[0] + sum(b * x for b, x in zip(beta[1:], logs))


def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    if group not in _COEFS:
        raise ValueError(f"Unknown group '{group}'. Available groups: {sorted(_COEFS.keys())}")

    group_coefs = _COEFS[group]
    outputs: List[Dict[str, float]] = []
    for point in input_data:
        pred: Dict[str, float] = {}
        for k in range(1, 6):
            key = f"loss_domain_{k}"
            beta = group_coefs[key]
            pred[key] = _predict_one(point, beta)
        outputs.append(pred)
    return outputs