# Project export: Mars Forecast This document was generated by HackStack to give an AI agent context about a hackathon project. Sections are labeled with their provenance; content marked as truncated was cut to keep this document small. ## Project metadata - Hackathon: UC Berkeley AI Hackathon 2026 - Tagline: Mars missions must forecast dangerous dust storms, but data is scarce. Proving an Earth-trained AI model generalizes to Mars unlocks cross-planetary transfer learning with minimal local data. - Devpost: https://devpost.com/software/mars-forecast - GitHub: https://github.com/LucasAschenbach/mars-weather - Demo: https://mars-weather-app-black.vercel.app/ - Video: https://www.youtube.com/embed/Vy_RPd0rblI?enablejsapi=1&hl=en_US&rel=0&start=&version=3&wmode=transparent - Team: 1 GitHub contributor(s) — Lucas Aschenbach (1 commits) ## Devpost submission (written by the team) No Devpost description available. ## README (from the GitHub repository) # Mars Weather Transfer Learning Can an Earth weather foundation model help forecast weather on Mars? This hackathon project adapts [Aurora](https://github.com/microsoft/aurora), a pretrained atmospheric model, to OpenMARS reanalysis data. The goal was to test whether Earth-weather pretraining transfers to a low-data planetary forecasting setting after fine-tuning on Mars observations. ![Mars forecast app](assets/mars-forecast-app.png) ## What We Built The repository contains two pieces: - `mars_weather/`: PyTorch data adapters, training utilities, and model construction for fine-tuning Aurora on OpenMARS. - `mars_weather_app/`: a Next.js dashboard for exploring exported Martian weather forecasts on a global Mars map. The model predicts global weather fields every 2 Martian hours on the OpenMARS latitude-longitude grid. Surface fields are predicted over `(lat, lon)`, and atmospheric fields are predicted over `(sigma level, lat, lon)`. | Field | Meaning | Units | Dimensions | | --- | --- | --- | --- | | `ps` | Surface pressure | Pa | `lat, lon` | | `tsurf` | Surface temperature | K | `lat, lon` | | `co2ice` | Surface CO2 ice mass | kg/m2 | `lat, lon` | | `dustcol` | Dust column opacity | opacity | `lat, lon` | | `u` | Zonal wind | m/s | `lev, lat, lon` | | `v` | Meridional wind | m/s | `lev, lat, lon` | | `temp` | Atmospheric temperature | K | `lev, lat, lon` | The OpenMARS vertical coordinate uses 35 sigma levels from near-surface levels to the upper atmosphere. ## Experiment Mars weather prediction is operationally important for landers, rovers, orbiters, entry-descent-landing planning, dust-storm risk, solar-power forecasting, and future crewed surface operations. The hard part is data scarcity: Earth weather models benefit from enormous observational and reanalysis archives, while Mars has a much smaller record. We fine-tuned Aurora on OpenMARS Mars Years 28-34 and validated on Mars Year 35. Samples are split by the Mars Year of the target frame using `splits/openmars_my28-34_train_my35_val.json`. We ran two main training configurations: - **Pretrained:** Aurora initialized from Earth-weather pretraining, then adapted to Mars variables and sigma levels. - **No pretraining:** the same Mars task setup, but without loading the pretrained Aurora checkpoint. Both models were evaluated with 20-step autoregressive rollouts. At 2 Martian hours per step, this corresponds to a 40-Martian-hour forecast horizon. Evaluation reports normalized MSE plus physical RMSE and MAE for each forecast field. ![20-step rollout RMSE comparison](assets/rollout_eval_comparison_rmse_grid.png) ## Result Across the 20-step rollout, the pretrained and non-pretrained models performed similarly on most aggregate metrics. The clearest difference was surface temperature: the Earth-pretrained model reduced `tsurf` prediction error by about 30% compared with the non-pretrained model. That result is preliminary because the hackathon timeline only allowed two main runs. Still, it is a useful signal: some of the structure learned from Earth weather appears to transfer to Mars after fine-tuning, even across different atmospheric composition, pressure regime, dust dynamics, radiative forcing, and planetary day length. ## Model Weights Model-only weights are published on Hugging Face as `safetensors` exports. These uploads strip optimizer state and RNG state from the original training checkpoints, reducing each artifact from about 15 GB to about 4.7 GB. | Model | Hugging Face repo | Notes | | --- | --- | --- | | Pretrained Aurora fine-tune | [`LucasAschenbach/mars-aurora-openmars-pretrained`](https://huggingface.co/LucasAschenbach/mars-aurora-openmars-pretrained) | Earth-pretrained Aurora base model fine-tuned on OpenMARS. | | Random-init comparison | [`LucasAschenbach/mars-aurora-openmars-random-init`](https://huggingface.co/LucasAschenbach/mars-aurora-openmars-random-init) | Same Aurora architecture trained on OpenMARS without loading the Earth-pretrained checkpoint. | Each model repository includes: - `model.safetensors`: model-only PyTorch state dict. - `openmars_stats.json`: normalization statistics needed to construct the Mars/OpenMARS model. - `config.json`: export metadata, including source checkpoint step. - `training_config.json` and `training_metadata.json`: run configuration and environment metadata. - `rollout_eval_val_20step.json`: 20-step recursive rollout metrics on the MY35 validation split. To reproduce the model-only export from a training checkpoint: ```bash python scripts/export_hf_weights.py \ --checkpoint artifacts/openmars_runs//checkpoint_step_9158.pt \ --run-dir artifacts/openmars_runs/ \ --output-dir artifacts/hf_exports/ \ --format safetensors \ --model-name ``` Upload the resulting folder with: ```bash hf upload-large-folder LucasAschenbach/ \ artifacts/hf_exports/ \ --repo-type model ``` ## Repository Layout ```text . ├── assets/ # README graphics ├── mars_weather/ # OpenMARS dataset, Aurora model, training utilities ├── mars_weather_app/ # Interactive forecast dashboard ├── scripts/ │ ├── create_openmars_split.py # Build reproducible train/validation split manifests │ ├── download_openmars.py # Download OpenMARS NetCDF files from Figshare │ ├── export_hf_weights.py # Strip training state and export HF-ready weights │ ├── export_rollout_netcdf.py # Export recursive rollouts to NetCDF │ ├── finetune_openmars.py # Fine-tune Aurora on OpenMARS │ ├── evaluate_openmars.py # Recursive rollout evaluation │ └── plot_rollout_eval.py # Plot rollout metrics ├── splits/ │ └── openmars_my28-34_train_my35_val.json └── tests/ ``` ## Setup Clone the repository with submodules so the local Aurora dependency is available: ```bash git clone --recurse-submodules https://github.com/LucasAschenbach/mars-weather cd mars-weather ``` Install Python dependencies: ```bash python -m venv .venv source .venv/bin/activate pip install -r requirements.txt ``` OpenMARS NetCDF files are expected under `data/`, matching the paths in the split manifest. To download the files used by the included MY28-34/MY35 split: ```bash python scripts/download_openmars.py \ --manifest splits/openmars_my28-34_train_my35_val.json \ --output-dir data \ --workers 4 ``` The downloader uses the Figshare file IDs already present in the split manifest, writes atomic `.part` files, resumes partial downloads when possible, retries failed transfers, and can optionally validate each NetCDF with `--validate`. ## Training Fine-tune with the provided Mars-Year split: ```bash python scripts/finetune_openmars.py \ --split-manifest splits/openmars_my28-34_train_my35_val.json \ --split train \ --model-size base \ --epochs 1 \ --batch-size 1 ``` To run the no-pretraining comparison, add: ```bash --no-load-checkpoint ``` The training script writes run configs, metadata, checkpoints, and TensorBoard logs under `artifacts/openmars_runs/` by default. To export a training checkpoint as model-only Hugging Face weights: ```bash python scripts/export_hf_weights.py \ --checkpoint artifacts/openmars_runs//checkpoint_latest.pt \ --run-dir artifacts/openmars_runs/ \ --output-dir artifacts/hf_exports/ \ --format safetensors \ --model-name ``` ## Evaluation Evaluate a trained checkpoint with a 20-step recursive rollout: ```bash python scripts/evaluate_openmars.py \ --run-dir artifacts/openmars_runs/ \ --split val \ --rollout-steps 20 \ --write-csv ``` Plot one or more rollout evaluation files: ```bash python scripts/plot_rollout_eval.py \ artifacts/openmars_runs//rollout_eval_val_20step.json \ artifacts/openmars_runs//rollout_eval_val_20step.json \ --labels pretrained scratch \ --o [README truncated for size] ## Detected evidence (automated analysis) Indexed codebase: 36 recognized source files, 382 KB. - CSS (language) — detected in the code - Next.js (technology) — detected in the code - Python (language) — detected in the code - PyTorch (technology) — detected in the code - React (technology) — detected in the code - Tailwind CSS (technology) — detected in the code - TypeScript (language) — detected in the code ## Codebase structure (from repository index) ### Files (120 of 392) ``` .gitignore .gitmodules mars_weather_app/.gitignore mars_weather_app/.vercelignore mars_weather_app/app/globals.css mars_weather_app/app/layout.tsx mars_weather_app/app/page.tsx mars_weather_app/components.json mars_weather_app/components/forecast-header.tsx mars_weather_app/components/forecast-source-picker.tsx mars_weather_app/components/layer-controls.tsx mars_weather_app/components/lead-time-slider.tsx mars_weather_app/components/mars-dashboard.tsx mars_weather_app/components/mars-map.tsx mars_weather_app/components/metric-card.tsx mars_weather_app/components/ui/button.tsx mars_weather_app/components/ui/popover.tsx mars_weather_app/components/vertical-level-selector.tsx mars_weather_app/eslint.config.mjs mars_weather_app/FORECAST_DATA.md mars_weather_app/lib/forecast-file.ts mars_weather_app/lib/mars-data.ts mars_weather_app/lib/utils.ts mars_weather_app/next-env.d.ts mars_weather_app/next.config.mjs mars_weather_app/package.json mars_weather_app/postcss.config.mjs mars_weather_app/public/forecasts/catalog.json mars_weather_app/public/forecasts/ground-truth/frames/lead-000.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-002.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-004.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-006.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-008.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-010.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-012.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-014.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-016.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-018.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-020.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-022.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-024.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-026.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-028.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-030.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-032.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-034.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-036.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-038.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-040.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-042.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-044.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-046.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-048.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-050.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-052.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-054.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-056.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-058.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-060.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-062.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-064.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-066.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-068.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-070.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-072.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-074.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-076.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-078.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-080.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-082.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-084.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-086.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-088.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-090.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-092.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-094.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-096.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-098.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-100.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-102.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-104.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-106.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-108.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-110.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-112.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-114.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-116.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-118.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-120.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-122.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-124.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-126.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-128.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-130.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-132.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-134.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-136.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-138.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-140.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-142.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-144.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-146.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-148.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-150.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-152.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-154.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-156.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-158.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-160.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-162.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-164.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-166.i16 mars_weather_app/public/forecasts/ground-truth/frames/lead-168.i16 mars_weather_app/public/forecasts/ground-truth/manifest.json mars_weather_app/public/forecasts/latest/frames/lead-000.i16 mars_weather_app/public/forecasts/latest/frames/lead-002.i16 mars_weather_app/public/forecasts/latest/frames/lead-004.i16 mars_weather_app/public/forecasts/latest/frames/lead-006.i16 mars_weather_app/public/forecasts/latest/frames/lead-008.i16 mars_weather_app/public/forecasts/latest/frames/lead-010.i16 [272 more files omitted for size] ``` ### Dependencies - mars_weather_app/package.json: @base-ui/react@^1.5.0, @tailwindcss/postcss@^4.2.0, @types/node@^24, @types/react@^19, @types/react-dom@^19, @vercel/analytics@1.6.1, class-variance-authority@^0.7.1, clsx@^2.1.1, eslint@^9.39.4, eslint-config-next@^16.2.9, lucide-react@^1.16.0, next@16.2.6, postcss@^8.5, react@^19, react-dom@^19, shadcn@^4.8.0, tailwind-merge@^3.3.1, tailwindcss@^4.2.0, tw-animate-css@^1.4.0, typescript@5.7.3 - requirements.txt: einops, h5netcdf, huggingface-hub, matplotlib, netcdf4, numpy, pydantic, pytest, scipy, tensorboard, timm, torch, tqdm, xarray ### Recent commits (newest first) - Add OpenMARS download script - Fix forecast frame selection consistency - Add selectable Mars forecast rollouts - Document Hugging Face model weights - Add Hugging Face weight export script - Add distributed OpenMARS rollout eval - Update clone command with actual repository URL - Document Mars weather project results - Add exported Mars forecast frames - Add Mars forecast dashboard - Add OpenMARS fine-tuning pipeline ## Key source files (fetched from GitHub, selected and truncated for size) ### mars_weather_app/FORECAST_DATA.md ```markdown # Forecast Data Keep raw model outputs out of `public/` and publish browser-ready frames into `public/forecasts/`. The app can load either a single legacy forecast at `public/forecasts/latest/` or a multi-forecast catalog at `public/forecasts/catalog.json`. Recommended layout: ```text mas_weather_app/ data/ forecasts/ raw/ latest.nc public/ forecasts/ catalog.json pretrained/ manifest.json frames/ lead-000.i16 lead-002.i16 random-init/ manifest.json frames/ lead-000.i16 lead-002.i16 ground-truth/ manifest.json frames/ lead-000.i16 lead-002.i16 ``` Export app-readable files from NetCDF rollouts: ```bash python scripts/export_forecast.py ../artifacts/app_rollouts/pretrained_7day_rollout.nc \ --output public/forecasts/pretrained \ --max-lead-hours 168 \ --source "Pretrained Aurora" ``` The catalog lists each selectable forecast: ```json { "schema": "mars-weather-forecast-catalog-v1", "forecasts": [ { "id": "pretrained", "label": "Pretrained", "manifestPath": "/forecasts/pretrained/manifest.json" } ] } ``` The app first tries to load: ```text /forecasts/catalog.json ``` If the catalog is missing, it falls back to the legacy single-forecast path: ```text /forecasts/latest/manifest.json ``` If no forecast manifest or frame can be loaded, it falls back to the deterministic preview data. The exported binary frame order is: ```text ps, tsurf, co2ice, dustcol, u, v, temp ``` Surface variables are stored as `lat, lon`; atmospheric variables are stored as `lev, lat, lon`. Values are quantized as signed 16-bit integers. The manifest stores each variable's `min` and `scale` for decoding in the browser. For aligned model-vs-truth comparisons, the exported manifests should share the same `baseTime`, `solarLongitude`, `marsYear`, and lead-hour frame sequence. The current 7-day app rollouts use 85 frames: lead 0 plus +2h through +168h. ``` ### requirements.txt ``` torch tensorboard numpy scipy matplotlib tqdm xarray netcdf4 h5netcdf einops timm huggingface-hub pydantic pytest ``` ### mars_weather_app/package.json ``` { "name": "martian-weather-forecast", "version": "0.1.0", "private": true, "scripts": { "dev": "next dev", "build": "next build", "start": "next start", "lint": "eslint ." }, "dependencies": { "@base-ui/react": "^1.5.0", "@vercel/analytics": "1.6.1", "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", "lucide-react": "^1.16.0", "next": "16.2.6", "react": "^19", "react-dom": "^19", "shadcn": "^4.8.0", "tailwind-merge": "^3.3.1", "tw-animate-css": "^1.4.0" }, "devDependencies": { "@tailwindcss/postcss": "^4.2.0", "@types/node": "^24", "@types/react": "^19", "@types/react-dom": "^19", "eslint": "^9.39.4", "eslint-config-next": "^16.2.9", "postcss": "^8.5", "tailwindcss": "^4.2.0", "typescript": "5.7.3" }, "pnpm": { "overrides": { "hono": "4.12.25" } }, "packageManager": "pnpm@9.12.3+sha512.cce0f9de9c5a7c95bef944169cc5dfe8741abfb145078c0d508b868056848a87c81e626246cb60967cbd7fd29a6c062ef73ff840d96b3c86c40ac92cf4a813ee" } ``` ### mars_weather_app/app/page.tsx ```typescript import { MarsDashboard } from "@/components/mars-dashboard" export default function Page() { return } ``` ### mars_weather_app/app/layout.tsx ```typescript import { Analytics } from '@vercel/analytics/next' import type { Metadata, Viewport } from 'next' import { Geist, Geist_Mono, Space_Grotesk } from 'next/font/google' import './globals.css' const geistSans = Geist({ variable: '--font-geist-sans', subsets: ['latin'] }) const geistMono = Geist_Mono({ variable: '--font-geist-mono', subsets: ['latin'], }) const spaceGrotesk = Space_Grotesk({ variable: '--font-space-grotesk', subsets: ['latin'], }) export const metadata: Metadata = { title: 'Martian Weather Forecast', description: 'OpenMARS Aurora weather prediction for Mars', generator: 'v0.app', icons: { icon: [ { url: '/icon-light-32x32.png', media: '(prefers-color-scheme: light)', }, { url: '/icon-dark-32x32.png', media: '(prefers-color-scheme: dark)', }, { url: '/icon.svg', type: 'image/svg+xml', }, ], apple: '/apple-icon.png', }, } export const viewport: Viewport = { colorScheme: 'dark', themeColor: '#2a1810', } export default function RootLayout({ children, }: Readonly<{ children: React.ReactNode }>) { return ( {children} {process.env.NODE_ENV === 'production' && } ) } ``` ### mars_weather_app/next-env.d.ts ```typescript /// /// import "./.next/types/routes.d.ts"; // NOTE: This file should not be edited // see https://nextjs.org/docs/app/api-reference/config/typescript for more information. ``` ### tests/conftest.py ```python """Test path setup for the nested Aurora checkout.""" from __future__ import annotations import sys from pathlib import Path ROOT = Path(__file__).resolve().parents[1] AURORA_ROOT = ROOT / "aurora" for path in (ROOT, AURORA_ROOT): path_str = str(path) if path_str not in sys.path: sys.path.insert(0, path_str) ``` ### mars_weather/_paths.py ```python """Import path helpers for the nested Aurora checkout.""" from __future__ import annotations import sys from pathlib import Path def ensure_aurora_on_path() -> None: """Make the vendored Aurora package importable from repo-root scripts.""" repo_root = Path(__file__).resolve().parents[1] aurora_root = repo_root / "aurora" if (aurora_root / "aurora" / "__init__.py").exists(): path = str(aurora_root) if path not in sys.path: sys.path.insert(0, path) ``` ### mars_weather/model.py ```python """Mars-specific Aurora model construction.""" from __future__ import annotations from datetime import timedelta from typing import Any, Literal from mars_weather._paths import ensure_aurora_on_path ensure_aurora_on_path() from aurora import AuroraPretrained, AuroraSmallPretrained from mars_weather.openmars import ( ATMOS_VARS, MARS_STATIC_VARS, OPENMARS_STEP, SURF_VARS, ) def make_mars_aurora( *, level_ids: tuple[int, ...], size: Literal["base", "small"] = "base", load_checkpoint: bool = True, timestep: timedelta = OPENMARS_STEP, use_lora: bool = False, autocast: bool = True, model_kwargs: dict[str, Any] | None = None, ): """Create an Aurora model configured for OpenMARS variables.""" kwargs: dict[str, Any] = { "surf_vars": SURF_VARS, "static_vars": MARS_STATIC_VARS, "atmos_vars": ATMOS_VARS, "level_condition": level_ids, "timestep": timestep, "use_lora": use_lora, "autocast": autocast, "positive_surf_vars": ("ps", "co2ice", "dustcol"), } if model_kwargs: kwargs.update(model_kwargs) cls = AuroraSmallPretrained if size == "small" else AuroraPretrained model = cls(**kwargs) if load_checkpoint: model.load_checkpoint(strict=False) return model ``` ### mars_weather/__init__.py ```python """Mars weather fine-tuning helpers for Aurora.""" from mars_weather._paths import ensure_aurora_on_path ensure_aurora_on_path() from mars_weather.model import make_mars_aurora from mars_weather.openmars import ( ATMOS_VARS, MARS_STATIC_VARS, SURF_VARS, OpenMARSDataset, OpenMARSRolloutDataset, OpenMARSStats, collate_batch_pairs, collate_rollout_samples, compute_openmars_stats, load_openmars_stats, register_openmars_stats, save_openmars_stats, ) from mars_weather.splits import ( create_openmars_split_manifest, dataset_from_manifest, load_split_manifest, rollout_dataset_from_manifest, save_split_manifest, ) from mars_weather.training import make_optimizer, normalized_mse_loss, normalized_mse_losses __all__ = [ "ATMOS_VARS", "MARS_STATIC_VARS", "SURF_VARS", "OpenMARSDataset", "OpenMARSRolloutDataset", "OpenMARSStats", "collate_batch_pairs", "collate_rollout_samples", "compute_openmars_stats", "create_openmars_split_manifest", "dataset_from_manifest", "load_openmars_stats", "load_split_manifest", "rollout_dataset_from_manifest", "make_mars_aurora", "make_optimizer", "normalized_mse_loss", "normalized_mse_losses", "register_openmars_stats", "save_split_manifest", "save_openmars_stats", ] ``` [26 more indexed source files omitted to keep this export small. The full file list is in the Codebase structure section above.]