import numpy as np
from healpix_plot.mollview import mollview

t_map = np.random.randn(12*64**2)

%load_ext autoreload
%autoreload 2

from healpix_analyse.healpix_sht import HEALPixSHT
import healpy as hp
import matplotlib.pyplot as plt

sht = HEALPixSHT(level=6)           # ~0.5 s de précomputation

# Spin-0 (scalaire)
alm  = sht.map2alm(t_map)           # différentiable
alm_hp = hp.map2alm(t_map,iter=0)
t_rec = sht.alm2map(alm)
t_rec_hp = hp.alm2map(alm_hp,nside=64)
cl   = sht.anafast(t_map)
cl_hp = hp.anafast(t_map,iter=0)

plt.figure(figsize=(12,4))
plt.subplot(1,4,1)
plt.plot(cl_hp,color='b',lw=4)
plt.plot(cl.cpu().numpy(),color='r')
plt.subplot(1,4,2)
plt.plot(cl_hp-cl.cpu().numpy())
plt.subplot(1,4,3)
plt.plot(abs(alm_hp))
plt.plot(abs(alm_hp-alm.cpu().numpy()))
plt.yscale('log')
plt.subplot(1,4,4)
plt.plot(abs(t_rec_hp))
plt.plot(abs(t_rec_hp-t_rec.cpu().numpy()))
plt.yscale('log')

import torch

nside=8
sht = HEALPixSHT(level=int(np.log2(nside)))  # ~0.5 s de précomputation
Q=np.random.randn(12*nside**2)
U=np.random.randn(12*nside**2)

plt.figure(figsize=(16,10))
for spin in range(3):
    # Spin-2 (Q/U → E/B) — pour l'apprentissage
    almE, almB = sht.map2alm_spin(Q,U, spin=spin)
    almE_hp, almB_hp = hp.map2alm_spin(np.vstack([Q,U]),spin)
    
    plt.subplot(3,4,1+4*spin)
    plt.title('E Real spin=%d'%(spin))
    plt.plot(almE_hp.real,color='b',lw=4,label='healpy')
    plt.plot(almE.cpu().numpy().real,color='r',label='healpix-analyse')
    plt.legend()
    plt.subplot(3,4,2+4*spin)
    plt.title('E imag spin=%d'%(spin))
    plt.plot(almE_hp.imag,color='b',lw=4)
    plt.plot(almE.cpu().numpy().imag,color='r')
    plt.subplot(3,4,3+4*spin)
    plt.title('B Real spin=%d'%(spin))
    plt.plot(almB_hp.real,color='b',lw=4)
    plt.plot(almB.cpu().numpy().real,color='r')
    plt.subplot(3,4,4+4*spin)
    plt.title('B imag spin=%d'%(spin))
    plt.plot(almB_hp.imag,color='b',lw=4)
    plt.plot(almB.cpu().numpy().imag,color='r')

plt.figure(figsize=(16,10))
for spin in range(3):
    # Spin-2 (Q/U → E/B) — pour l'apprentissage
    almE, almB = sht.map2alm_spin(Q,U, spin=spin)
    almE_hp, almB_hp = hp.map2alm_spin(np.vstack([Q,U]),spin)
    
    Q_rec, U_rec = sht.alm2map_spin(almE, almB, spin=spin)
    Q_rec_hp, U_rec_hp = hp.alm2map_spin([almE_hp, almB_hp],nside, spin,3*nside-1)
    
    plt.subplot(2,3,1+spin)
    plt.title('Q spin=%d'%(spin))
    plt.plot(Q_rec_hp.real,color='b',lw=4,label='healpy')
    plt.plot(Q_rec.cpu().numpy().real,color='r',label='healpix-analyse')
    plt.legend()
    plt.subplot(2,3,4+spin)
    plt.title('U spin=%d'%(spin))
    plt.plot(U_rec_hp.real,color='b',lw=4,label='healpy')
    plt.plot(U_rec.cpu().numpy().real,color='r',label='healpix-analyse')

# Curl/Div d'un champ vectoriel (spin=1)
div, curl = sht.uv_to_curl_div(Q, U)

try:
    import cdsapi
except:
    !pip install xarray cdsapi

#!/usr/bin/env python
import xarray as xr

try:
    dsG = xr.open_dataset("outputGT10m.grib")
except:
    c = cdsapi.Client()
    
    c.retrieve("reanalysis-era5-complete", {
        "class": "ea",
        "date": "2025-01-01",
        "expver": "1",
        "levtype": "sfc",
        "number": "0/1/2/3/4/5/6/7/8/9",
        "param": "165.128/166.128", #u,v at 10m
        "step": "12",
        "stream": "enda",
        "time": "09:00:00",
        "type": "4v",
        "grid":"N256",
    }, "outputGT10m.grib")
    dsG = xr.open_dataset("outputGT10m.grib")

dsG

/srv/conda/envs/notebook/lib/python3.13/site-packages/xarray/backends/plugins.py:109: RuntimeWarning: Engine 'kerchunk' loading failed:
No module named 'zarr.core.array_spec'; 'zarr.core' is not a package
  external_backend_entrypoints = backends_dict_from_pkg(entrypoints_unique)
Ignoring index file 'outputGT10m.grib.5b7b6.idx' older than GRIB file

<xarray.Dataset> Size: 47MB
Dimensions:     (number: 10, values: 348528)
Coordinates:
  * number      (number) int64 80B 0 1 2 3 4 5 6 7 8 9
    latitude    (values) float64 3MB ...
    longitude   (values) float64 3MB ...
    time        datetime64[ns] 8B ...
    step        timedelta64[ns] 8B ...
    surface     float64 8B ...
    valid_time  datetime64[ns] 8B ...
Dimensions without coordinates: values
Data variables:
    u10         (number, values) float32 14MB ...
    v10         (number, values) float32 14MB ...
    sst         (number, values) float32 14MB ...
Attributes:
    GRIB_edition:            1
    GRIB_centre:             ecmf
    GRIB_centreDescription:  European Centre for Medium-Range Weather Forecasts
    GRIB_subCentre:          0
    Conventions:             CF-1.7
    institution:             European Centre for Medium-Range Weather Forecasts
    history:                 2026-08-03T14:12 GRIB to CDM+CF via cfgrib-0.9.1...

xarray.Dataset

• Dimensions:
  • number: 10
  • values: 348528
• Coordinates: (7)
  • number
    (number)
    int64
    0 1 2 3 4 5 6 7 8 9

    long_name :

    ensemble member numerical id

    units :

    1

    standard_name :

    realization

    array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

  • latitude
    (values)
    float64
    ...

    units :

    degrees_north

    standard_name :

    latitude

    long_name :

    latitude

    [348528 values with dtype=float64]

  • longitude
    (values)
    float64
    ...

    units :

    degrees_east

    standard_name :

    longitude

    long_name :

    longitude

    [348528 values with dtype=float64]

  • time
    ()
    datetime64[ns]
    ...

    long_name :

    initial time of forecast

    standard_name :

    forecast_reference_time

    [1 values with dtype=datetime64[ns]]

  • step
    ()
    timedelta64[ns]
    ...

    long_name :

    time since forecast_reference_time

    standard_name :

    forecast_period

    [1 values with dtype=timedelta64[ns]]

  • surface
    ()
    float64
    ...

    long_name :

    original GRIB coordinate for key: level(surface)

    units :

    1

    [1 values with dtype=float64]

  • valid_time
    ()
    datetime64[ns]
    ...

    standard_name :

    time

    long_name :

    time

    [1 values with dtype=datetime64[ns]]

• Data variables: (3)
  • u10
    (number, values)
    float32
    ...

    GRIB_paramId :

    165

    GRIB_dataType :

    4v

    GRIB_numberOfPoints :

    348528

    GRIB_typeOfLevel :

    surface

    GRIB_stepUnits :

    1

    GRIB_stepType :

    instant

    GRIB_gridType :

    reduced_gg

    GRIB_uvRelativeToGrid :

    0

    GRIB_N :

    256

    GRIB_NV :

    0

    GRIB_cfName :

    unknown

    GRIB_cfVarName :

    u10

    GRIB_gridDefinitionDescription :

    Gaussian Latitude/Longitude Grid

    GRIB_missingValue :

    3.4028234663852886e+38

    GRIB_name :

    10 metre U wind component

    GRIB_pl :

    [ 18 25 32 40 45 50 60 64 72 72 75 81 90 96 100 108 120 120 125 135 144 150 160 160 180 180 180 192 192 200 216 216 216 225 240 240 243 250 256 270 270 288 288 288 300 300 320 320 320 324 360 360 360 360 360 360 375 375 384 384 400 400 400 432 432 432 432 432 450 450 450 480 480 480 480 480 486 500 500 500 512 512 540 540 540 540 540 576 576 576 576 576 576 600 600 600 600 600 640 640 640 640 640 640 640 640 648 675 675 675 675 675 675 720 720 720 720 720 720 720 720 720 729 729 750 750 750 750 750 768 768 768 768 800 800 800 800 800 800 800 800 810 810 864 864 864 864 864 864 864 864 864 864 864 864 864 864 900 900 900 900 900 900 900 900 900 900 900 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 972 972 972 972 972 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 972 972 972 972 972 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 900 900 900 900 900 900 900 900 900 900 900 864 864 864 864 864 864 864 864 864 864 864 864 864 864 810 810 800 800 800 800 800 800 800 800 768 768 768 768 750 750 750 750 750 729 729 720 720 720 720 720 720 720 720 720 675 675 675 675 675 675 648 640 640 640 640 640 640 640 640 600 600 600 600 600 576 576 576 576 576 576 540 540 540 540 540 512 512 500 500 500 486 480 480 480 480 480 450 450 450 432 432 432 432 432 400 400 400 384 384 375 375 360 360 360 360 360 360 324 320 320 320 300 300 288 288 288 270 270 256 250 243 240 240 225 216 216 216 200 192 192 180 180 180 160 160 150 144 135 125 120 120 108 100 96 90 81 75 72 72 64 60 50 45 40 32 25 18]

    GRIB_shortName :

    10u

    GRIB_totalNumber :

    10

    GRIB_units :

    m s**-1

    long_name :

    10 metre U wind component

    units :

    m s**-1

    standard_name :

    unknown

    [3485280 values with dtype=float32]

  • v10
    (number, values)
    float32
    ...

    GRIB_paramId :

    166

    GRIB_dataType :

    4v

    GRIB_numberOfPoints :

    348528

    GRIB_typeOfLevel :

    surface

    GRIB_stepUnits :

    1

    GRIB_stepType :

    instant

    GRIB_gridType :

    reduced_gg

    GRIB_uvRelativeToGrid :

    0

    GRIB_N :

    256

    GRIB_NV :

    0

    GRIB_cfName :

    unknown

    GRIB_cfVarName :

    v10

    GRIB_gridDefinitionDescription :

    Gaussian Latitude/Longitude Grid

    GRIB_missingValue :

    3.4028234663852886e+38

    GRIB_name :

    10 metre V wind component

    GRIB_pl :

    [ 18 25 32 40 45 50 60 64 72 72 75 81 90 96 100 108 120 120 125 135 144 150 160 160 180 180 180 192 192 200 216 216 216 225 240 240 243 250 256 270 270 288 288 288 300 300 320 320 320 324 360 360 360 360 360 360 375 375 384 384 400 400 400 432 432 432 432 432 450 450 450 480 480 480 480 480 486 500 500 500 512 512 540 540 540 540 540 576 576 576 576 576 576 600 600 600 600 600 640 640 640 640 640 640 640 640 648 675 675 675 675 675 675 720 720 720 720 720 720 720 720 720 729 729 750 750 750 750 750 768 768 768 768 800 800 800 800 800 800 800 800 810 810 864 864 864 864 864 864 864 864 864 864 864 864 864 864 900 900 900 900 900 900 900 900 900 900 900 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 972 972 972 972 972 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 972 972 972 972 972 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 900 900 900 900 900 900 900 900 900 900 900 864 864 864 864 864 864 864 864 864 864 864 864 864 864 810 810 800 800 800 800 800 800 800 800 768 768 768 768 750 750 750 750 750 729 729 720 720 720 720 720 720 720 720 720 675 675 675 675 675 675 648 640 640 640 640 640 640 640 640 600 600 600 600 600 576 576 576 576 576 576 540 540 540 540 540 512 512 500 500 500 486 480 480 480 480 480 450 450 450 432 432 432 432 432 400 400 400 384 384 375 375 360 360 360 360 360 360 324 320 320 320 300 300 288 288 288 270 270 256 250 243 240 240 225 216 216 216 200 192 192 180 180 180 160 160 150 144 135 125 120 120 108 100 96 90 81 75 72 72 64 60 50 45 40 32 25 18]

    GRIB_shortName :

    10v

    GRIB_totalNumber :

    10

    GRIB_units :

    m s**-1

    long_name :

    10 metre V wind component

    units :

    m s**-1

    standard_name :

    unknown

    [3485280 values with dtype=float32]

  • sst
    (number, values)
    float32
    ...

    GRIB_paramId :

    34

    GRIB_dataType :

    4v

    GRIB_numberOfPoints :

    348528

    GRIB_typeOfLevel :

    surface

    GRIB_stepUnits :

    1

    GRIB_stepType :

    instant

    GRIB_gridType :

    reduced_gg

    GRIB_uvRelativeToGrid :

    0

    GRIB_N :

    256

    GRIB_NV :

    0

    GRIB_cfName :

    unknown

    GRIB_cfVarName :

    sst

    GRIB_gridDefinitionDescription :

    Gaussian Latitude/Longitude Grid

    GRIB_missingValue :

    3.4028234663852886e+38

    GRIB_name :

    Sea surface temperature

    GRIB_pl :

    [ 18 25 32 40 45 50 60 64 72 72 75 81 90 96 100 108 120 120 125 135 144 150 160 160 180 180 180 192 192 200 216 216 216 225 240 240 243 250 256 270 270 288 288 288 300 300 320 320 320 324 360 360 360 360 360 360 375 375 384 384 400 400 400 432 432 432 432 432 450 450 450 480 480 480 480 480 486 500 500 500 512 512 540 540 540 540 540 576 576 576 576 576 576 600 600 600 600 600 640 640 640 640 640 640 640 640 648 675 675 675 675 675 675 720 720 720 720 720 720 720 720 720 729 729 750 750 750 750 750 768 768 768 768 800 800 800 800 800 800 800 800 810 810 864 864 864 864 864 864 864 864 864 864 864 864 864 864 900 900 900 900 900 900 900 900 900 900 900 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 972 972 972 972 972 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1024 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 972 972 972 972 972 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 960 900 900 900 900 900 900 900 900 900 900 900 864 864 864 864 864 864 864 864 864 864 864 864 864 864 810 810 800 800 800 800 800 800 800 800 768 768 768 768 750 750 750 750 750 729 729 720 720 720 720 720 720 720 720 720 675 675 675 675 675 675 648 640 640 640 640 640 640 640 640 600 600 600 600 600 576 576 576 576 576 576 540 540 540 540 540 512 512 500 500 500 486 480 480 480 480 480 450 450 450 432 432 432 432 432 400 400 400 384 384 375 375 360 360 360 360 360 360 324 320 320 320 300 300 288 288 288 270 270 256 250 243 240 240 225 216 216 216 200 192 192 180 180 180 160 160 150 144 135 125 120 120 108 100 96 90 81 75 72 72 64 60 50 45 40 32 25 18]

    GRIB_shortName :

    sst

    GRIB_totalNumber :

    10

    GRIB_units :

    K

    long_name :

    Sea surface temperature

    units :

    K

    standard_name :

    unknown

    [3485280 values with dtype=float32]

• Attributes: (7)

  GRIB_edition :

  1

  GRIB_centre :

  ecmf

  GRIB_centreDescription :

  European Centre for Medium-Range Weather Forecasts

  GRIB_subCentre :

  0

  Conventions :

  CF-1.7

  institution :

  European Centre for Medium-Range Weather Forecasts

  history :

  2026-08-03T14:12 GRIB to CDM+CF via cfgrib-0.9.15.1/ecCodes-2.47.0 with {"source": "outputGT10m.grib", "filter_by_keys": {}, "encode_cf": ["parameter", "time", "geography", "vertical"]}

convert ERA5 data into healpix using ALM

from healpix_analyse.alm_latlon import build_rings_from_latlon,map2alm_latlon

ring_theta, ring_phi_list, ring_counts, idx = build_rings_from_latlon(np.deg2rad(90.0-dsG.latitude.values),
                         np.deg2rad(dsG.longitude.values))

nside_out=64 # keep small to get rapid results

almu = map2alm_latlon(torch.tensor(dsG.u10.values[0,idx],device='cuda:0'),
                    ring_theta, ring_phi_list, ring_counts,quadrature='gauss_legendre',lmax=nside_out*3-1)
almv = map2alm_latlon(torch.tensor(dsG.v10.values[0,idx],device='cuda:0'),
                    ring_theta, ring_phi_list, ring_counts,quadrature='gauss_legendre',lmax=nside_out*3-1)

sht = HEALPixSHT(level=int(np.log2(nside_out)))

imu = sht.alm2map(almu)
imv = sht.alm2map(almv)

comput div,curl from u,v field

div, curl = sht.uv_to_curl_div(imu, imv)

from healpix_plot.mollview import mollview

plt.figure(figsize=(16,10))
mollview(imu.cpu().numpy(),cmap='coolwarm',title=str(dsG.time.values)+' u10',hold=False,sub=(2,2,1))
mollview(imv.cpu().numpy(),cmap='coolwarm',title=str(dsG.time.values)+' v10',hold=False,sub=(2,2,2))
mollview(div.cpu().numpy(),cmap='coolwarm',title=str(dsG.time.values)+' div',hold=False,sub=(2,2,3))
mollview(curl.cpu().numpy(),cmap='coolwarm',title=str(dsG.time.values)+' curl',hold=False,sub=(2,2,4))