From 4b3b97eebd244b54843c294bfdc3904eb7ba7fd0 Mon Sep 17 00:00:00 2001 From: arzoo0511 Date: Tue, 31 Mar 2026 14:41:56 +0530 Subject: [PATCH] feat: add spherical geometry utilities for global domain support --- .cora/analytics/index.json | 1 + geometry.py | 19 +++++++++++++++++++ neural_lam/geometry.py | 0 3 files changed, 20 insertions(+) create mode 100644 .cora/analytics/index.json create mode 100644 geometry.py create mode 100644 neural_lam/geometry.py diff --git a/.cora/analytics/index.json b/.cora/analytics/index.json new file mode 100644 index 000000000..0e637c9da --- /dev/null +++ b/.cora/analytics/index.json @@ -0,0 +1 @@ +{"version":2,"userOnlyFiles":{"tests\\__init__.py":1,".cirun.yml":17,"AGENTS.md":98,"neural_lam\\__init__.py":18,"neural_lam\\config.py":189,".pre-commit-config.yaml":45,"CHANGELOG.md":318,"neural_lam\\create_graph.py":614,"neural_lam\\custom_loggers.py":69,"neural_lam\\datastore\\__init__.py":28,"neural_lam\\datastore\\base.py":643,"neural_lam\\datastore\\npyfilesmeps\\__init__.py":3,"neural_lam\\datastore\\mdp.py":546,"neural_lam\\datastore\\npyfilesmeps\\compute_standardization_stats.py":428,"neural_lam\\datastore\\npyfilesmeps\\config.py":68,"neural_lam\\datastore\\npyfilesmeps\\store.py":809,"neural_lam\\datastore\\plot_example.py":191,"neural_lam\\interaction_net.py":172,"neural_lam\\loss_weighting.py":107,"neural_lam\\metrics.py":238,"neural_lam\\models\\__init__.py":7,"neural_lam\\models\\ar_model.py":825,"neural_lam\\models\\base_graph_model.py":366,"neural_lam\\models\\base_hi_graph_model.py":238,"neural_lam\\models\\graph_lam.py":92,"neural_lam\\models\\hi_lam_parallel.py":100,"neural_lam\\models\\hi_lam.py":208,"neural_lam\\train_model.py":383,"pyproject.toml":133,"neural_lam\\plot_graph.py":288,"neural_lam\\vis.py":330,"neural_lam\\utils.py":623,"tests\\datastore_examples\\mdp\\danra_100m_winds\\config.yaml":19,"neural_lam\\weather_dataset.py":741,"README.md":587,"tests\\conftest.py":126,"tests\\datastore_examples\\mdp\\danra_100m_winds\\danra.datastore.yaml":119,"tests\\datastore_examples\\mdp\\danra_100m_winds\\README.md":6,"tests\\datastore_examples\\npyfilesmeps\\README.md":6,"tests\\dummy_datastore.py":761,"tests\\test_clamping.py":284,"tests\\test_cli.py":120,"tests\\test_config.py":73,"tests\\test_datasets.py":499,"tests\\test_graph_creation.py":120,"tests\\test_imports.py":10,"tests\\test_datastores.py":444,"tests\\test_plotting.py":412,"tests\\test_time_slicing.py":196,"tests\\test_plot_graph.py":85,"tests\\test_train_model_warnings.py":44,"tests\\test_training.py":192,"geometry.py":19},"trackedFiles":[]} \ No newline at end of file diff --git a/geometry.py b/geometry.py new file mode 100644 index 000000000..de9d412cf --- /dev/null +++ b/geometry.py @@ -0,0 +1,19 @@ +import torch + +def lat_lon_to_cartesian(lat, lon): + lat_rad = torch.deg2rad(lat) + lon_rad = torch.deg2rad(lon) + + x = torch.cos(lat_rad) * torch.cos(lon_rad) + y = torch.cos(lat_rad) * torch.sin(lon_rad) + z = torch.sin(lat_rad) + + return torch.stack([x, y, z], dim=-1) + + +def get_area_weights(lat): + return torch.cos(torch.deg2rad(lat)) + +def get_cartesian_coords(self): + lat, lon = self.get_lat_lon() + return lat_lon_to_cartesian(lat, lon) \ No newline at end of file diff --git a/neural_lam/geometry.py b/neural_lam/geometry.py new file mode 100644 index 000000000..e69de29bb