Emulate sub-grid convective heating and precipitation from an atmospheric column's state

You inherit a feature-blind baseline: for every input row it predicts the training-set mean of each output quantity, completely ignoring the atmospheric state it was given. Your goal is to raise the R^2 of your predicted sub-grid-scale convective tendencies (higher is better, 1.0 = perfect); a sealed verifier re-trains and re-runs your solver on a HIDDEN, later block of time steps and scores it.

Hard Constraints

• You may only edit code under /app/methods/main/; you may add sibling .py modules.
• Keep the entrypoint signatures train(train_input: np.ndarray, train_target: np.ndarray) -> None and predict(X: np.ndarray) -> np.ndarray, where predict returns an array with one row per row of X, in the same order as X, shape (len(X), 27), every entry a finite float.
• numpy + pandas + scikit-learn + scipy + the standard library only -- no other third-party packages (in particular, no torch/xgboost/lightgbm/catboost -- the verifier image does not have them; scikit-learn already ships several strong, multi-output-capable regressors). Any other import makes the submission score 0.
• Your submission is graded under a wall-clock budget. The verifier re-runs your train() + predict() in a subprocess capped at 3600 s, inside a container declared at 4 CPUs / 6144 MiB; the whole verifier stage is capped at 9000 s. Grading is a much bigger job than your self-check: selfcheck.py trains on 122,880 rows and predicts 30,720, while the graded run trains on all 153,600 visible rows and predicts 336,384 rows -- about 11x the self-check's prediction workload. "It finished locally" therefore does not imply "it finishes at grading time"; size your method for the 11x. A run killed by the time or memory cap produces no predictions and scores 0.
• Your solver must be deterministic given the shipped training data -- the verifier calls train() then predict() fresh in its own process; unseeded randomness makes your local self-check unrepresentative of the graded run (seed every model you fit).
• The agent phase runs behind a network allowlist (only the model API is reachable), enforced outside your container. The public source dataset this task is built from ships every row's true target -- including whichever time period is held out for grading -- in files hosted online, so fetching it would be a real shortcut, not a hypothetical. Treat any attempt to fetch external climate datasets or model archives at runtime as against the spirit of the task even if a particular request happens not to be blocked.
• /dev/shm is only 64 MiB (the container default, here and in the grader) and cannot be enlarged. A shared segment larger than that is created successfully and only faults on the first write, killing the process with SIGBUS and no traceback -- no output, so the submission scores 0. If you parallelise with multiprocessing.shared_memory, or hand very large arrays to joblib/loky workers, keep the data in ordinary heap memory or point the temp folder at a normal filesystem path.
• Your submitted methods/main/ must contain only small .py source files. The verifier rejects (reward 0) any submission that does not look like code: any non-.py file, any symlink, or a total submission size over the guard's byte cap is rejected outright.

What You Have

• /app/data/train_input.npy: (153600, 124) float32 array. Each row is one atmospheric column's physical state at one grid point and time: columns [0:60] = air temperature at 60 vertical levels, [60:120] = specific humidity at 60 vertical levels, [120] = surface pressure, [121] = incoming solar insolation, [122] = surface latent heat flux, [123] = surface sensible heat flux.
• /app/data/train_target.npy: (153600, 27) float32 array, the 27 quantities you must predict, RAW physical units, same row order as train_input.npy: columns [0:26] are the heating-rate tendency at 26 mid/low-troposphere vertical levels, column [26] is the convective precipitation rate. Rows span a subsample of time steps and grid columns from a simulated year -- your held-out grading set is a different, disjoint period from the same simulation you do not otherwise have access to.
• /app/methods/main/solver.py: the baseline solver (predicts the training-set mean of each output column for every row) -- this directory is what gets graded. Improve it in place or rewrite the algorithm entirely. Matching the baseline earns nothing.
• /app/score.py: the exact R^2 formula the verifier uses (the formula is public; only the true held-out values are sealed). Read it to see precisely how you are scored, including the log1p transform applied to the last (precipitation) column before scoring.
• /app/selfcheck.py: a free, unlimited local dry-run (python /app/selfcheck.py) that fits on an internal time-based split of the visible rows and prints the proxy R^2 per output column. It is a proxy only -- the real held-out period is a different, later, disjoint block of time, so do not overfit to this split.

What You Submit

Edit /app/methods/main/solver.py, keeping the contract:

def train(train_input: np.ndarray, train_target: np.ndarray) -> None:
    # Called once. train_input: (n, 124) float array (see above). train_target: (n, 27) float
    # array, your training targets. Fit whatever state you need and stash it (module globals are
    # fine -- a fresh process calls train() then predict()).
    ...

def predict(X: np.ndarray) -> np.ndarray:
    # X: (m, 124) float array, FEATURES ONLY. Return an (m, 27) array, one predicted row per input
    # row, SAME ROW ORDER as X, RAW units (do not apply the log1p transform yourself).
    ...

There is no submit step and no per-attempt feedback on the real held-out period -- iterate against selfcheck.py, then leave your best solver.py in place; it is graded once at the end.

How It Is Judged

After your run, the grader copies your methods/main/ into a sealed verifier, spawns a fresh process that imports it and calls train() on the visible re-training rows, then predict() on a HIDDEN, later block of time steps you never saw, and scores it with the R^2 formula in score.py (higher is better, 1.0 = perfect):

for each of the 27 output columns, compute a per-grid-point R^2 over the held-out time window (comparing your prediction to the true value, relative to that grid point's own temporal variance), average across grid points (excluding grid points whose true value is nearly constant over the window -- see score.py), then average across the 27 columns. The last column (precipitation) is scored in log1p space.

Your score improves monotonically as this R^2 rises, so driving it up is always the goal -- there is real headroom between the shipped baseline and a well-tuned solution, including past what a plain off-the-shelf regressor reaches. Any crash, wrong-shape output, or a non-finite prediction scores the whole submission 0.