feat: add spherical geometry utilities for global domain support

.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":[]}

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)