%load_ext autoreload
%autoreload 2

import numpy as np
import matplotlib.pyplot as plt

from healpix_plot.mollview import mollview

im = np.load('./heal_im.npy')

from healpix_analyse.alm_latlon import build_rings_from_latlon,anafast_latlon
import healpix_geo as hg
import torch

#im=np.random.randn(12*nside**2)
#im=hp.ud_grade(im,nside)
ilevel=int(np.log2(im.shape[0]//12)/2)

p,t=hg.ring.healpix_to_lonlat(np.arange(12*4**ilevel),ilevel,ellipsoid='sphere')
ring_theta, ring_phi_list, ring_counts, idx = build_rings_from_latlon(np.deg2rad(90-t),np.deg2rad(p))

mollview(im,nest=True)

try:
    cl2.np.load('./heal_im_cl.npy')
except:
    import healpy as hp
    cl2 = hp.anafast(im)
    np.save('./heal_im_cl.npy',cl2)
    
cl = anafast_latlon(im[idx], ring_theta, ring_phi_list, ring_counts,lmax=3*2**ilevel-1)

plt.figure(figsize=(16,6))
plt.plot(cl2,lw=4,label='anafast healpy',color='b')
plt.plot(cl.cpu().numpy(),label='anafast_latlon',color='r')

plt.xscale('log')
plt.yscale('log')
plt.xlabel('Multipole')
plt.legend()
plt.xlabel('C(l)')

Text(0.5, 0, 'C(l)')

try:
    import cdsapi
except:
    !pip install xarray cdsapi

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

try:
    dsG = xr.open_dataset("outputGT2m.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": "167.128",
        "step": "12",
        "stream": "enda",
        "time": "09:00:00",
        "type": "4v",
        "grid":"N256",
    }, "outputGT2m.grib")
    dsG = xr.open_dataset("outputGT2m.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)

<xarray.Dataset> Size: 20MB
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:
    t2m         (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-03T13:04 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: (1)
  • t2m
    (number, values)
    float32
    ...

    GRIB_paramId :

    167

    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 :

    t2m

    GRIB_gridDefinitionDescription :

    Gaussian Latitude/Longitude Grid

    GRIB_missingValue :

    3.4028234663852886e+38

    GRIB_name :

    2 metre 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 :

    2t

    GRIB_totalNumber :

    10

    GRIB_units :

    K

    long_name :

    2 metre 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-03T13:04 GRIB to CDM+CF via cfgrib-0.9.15.1/ecCodes-2.47.0 with {"source": "outputGT2m.grib", "filter_by_keys": {}, "encode_cf": ["parameter", "time", "geography", "vertical"]}

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

%%time
cl = anafast_latlon(torch.tensor(dsG.t2m.values[0,idx],device='cpu'),
                    ring_theta, ring_phi_list, ring_counts,quadrature='gauss_legendre',lmax=256)

CPU times: user 4.19 s, sys: 386 ms, total: 4.57 s
Wall time: 621 ms

plt.figure(figsize=(16,6))
plt.plot(cl.cpu().numpy(),label='anafast_latlon',color='r')

plt.xscale('log')
plt.yscale('log')
plt.xlabel('Multipole')
plt.legend()
plt.xlabel('C(l)')

Text(0.5, 0, 'C(l)')

from healpix_analyse.convol import HealPixConv
from healpix_analyse.down   import HealPixDown
from healpix_analyse.up     import HealPixUp

import torch
nside = 64

down = HealPixDown(level=int(np.log2(nside)))
conv = HealPixConv(level=int(np.log2(nside)), in_channels=1, out_channels=16)
up   = HealPixUp(level=int(np.log2(nside)) - 1)

x = torch.randn(8, 12 * nside**2)       # [B, N]
x_feat = conv(x)                         # [B, 16, N]
x_down, ids_out = down(x_feat[:, 0, :]) # [B, N/4]
x_up, ids_fine  = up(x_down)            # [B, N]

from healpix_analyse.down   import HealPixDown
from healpix_analyse.up     import HealPixUp
nside=64
import numpy as np
import matplotlib.pyplot as plt

plt.figure(figsize=(12,3))
x=np.random.randn(12*nside**2)
x_down=x
mollview(x_down,nest=True,sub=(2,5,1),hold=False,title='Input data')
lnside=nside
for k in range(4):
    down = HealPixDown(level=int(np.log2(lnside)),ellipsoid='sphere')
    x_down, _ = down(x_down)
    lnside/=2
    mollview(x_down,nest=True,sub=(2,5,2+k),hold=False,title='down %d'%(k))

x_up=x_down
for k in range(4):
    up   = HealPixUp(level=int(np.log2(lnside)),ellipsoid='sphere')
    x_up, _  = up(x_up)
    lnside*=2
    mollview(x_up,nest=True,sub=(2,5,7+k),hold=False,title='up %d'%(k))

import numpy as np
from healpix_analyse.convol import HealPixConv
import matplotlib.pyplot as plt
nside = 64
kernelsz=5
conv = HealPixConv(level=int(np.log2(nside)), in_channels=1, out_channels=1, n_gauges=4, kernel_sz=kernelsz,gauge_type='phi')

# 25-point Gaussian: center gets more weight than the 8 neighbours
x,y = np.meshgrid(np.arange(kernelsz)-kernelsz//2,np.arange(kernelsz)-kernelsz//2)

W = np.zeros([1,1,kernelsz*kernelsz])
W[0,0] = (np.exp(-0.25*(x**2+y**2))*np.cos(0.5*(x)*np.pi)).ravel()
W[0,0] -= W[0,0].mean()

plt.imshow(W[0,0].reshape(kernelsz,kernelsz))
conv.set_kernel(W)

HealPixConv(level=6, nside=64, in=1, out=1, kernel_sz=5, P=25, G=4, gauge='phi', partial=False)

import numpy as np
sky = np.random.randn(12 * nside**2).astype(np.float32)
sky_smooth = conv(sky)     # returns np.ndarray [N]
print(sky_smooth.shape)
plt.figure(figsize=(16,9))
for k in range(4):
    mollview(sky_smooth[k],nest=True,sub=(2,2,1+k),hold=False,title=r'$\Theta=%d$'%(k))    

(4, 49152)

nside = 64
kernelsz=7
conv = HealPixConv(level=int(np.log2(nside)), in_channels=1, out_channels=1, n_gauges=4, kernel_sz=kernelsz,gauge_type='phi')

# 25-point Gaussian: center gets more weight than the 8 neighbours
W = np.array([[0,0,0,0,1,0,0],
              [0,0,0,0,1,1,0],
              [0,0,0,0,1,1,1],
              [1,1,1,1,1,1,1],
              [0,0,0,0,1,1,1],
              [0,0,0,0,1,1,0],
              [0,0,0,0,1,0,0]]).ravel().astype('float32').reshape(1,1,kernelsz*kernelsz)
    

plt.imshow(W[0,0].reshape(kernelsz,kernelsz))
conv.set_kernel(W)
sky = np.zeros([12 * nside**2],dtype=np.float32)
tc,pc=hp.pix2ang(4,np.arange(12*4**2),nest=True)
sky[hp.ang2pix(nside,tc,pc,nest=True)]=1.0
sky_smooth = conv(sky)     # returns np.ndarray [N]
print(sky_smooth.shape)
plt.figure(figsize=(16,9))
for k in range(4):
    mollview(sky_smooth[k],nest=True,sub=(2,2,1+k),hold=False,title=r'gauge="phi" $\Theta=%d$'%(k))    # (49152,)

(4, 49152)

nside = 64
kernelsz=7
conv = HealPixConv(level=int(np.log2(nside)), in_channels=1, out_channels=1,device='cuda:0',
                   n_gauges=4, kernel_sz=kernelsz,gauge_type='cosmo')

# 25-point Gaussian: center gets more weight than the 8 neighbours
W = np.array([[0,0,0,0,1,0,0],
              [0,0,0,0,1,1,0],
              [0,0,0,0,1,1,1],
              [1,1,1,1,1,1,1],
              [0,0,0,0,1,1,1],
              [0,0,0,0,1,1,0],
              [0,0,0,0,1,0,0]]).ravel().astype('float32').reshape(1,1,kernelsz*kernelsz)
conv.set_kernel(W)

sky = np.zeros([12 * nside**2],dtype=np.float32)
tc,pc=hp.pix2ang(4,np.arange(12*4**2),nest=True)
sky[hp.ang2pix(nside,tc,pc,nest=True)]=1.0
sky_smooth = conv(sky)     # returns np.ndarray [N]
print(sky_smooth.shape)
plt.figure(figsize=(16,10))
for k in range(4):
    mollview(sky_smooth[k],nest=True,sub=(2,2,1+k),hold=False,title=r'gauge="cosmo" $\Theta=%d$'%(k))    # (49152,)

(4, 49152)

# Placer les singularités loin de l'Europe (vers le Pacifique et l'océan Indien)
conv = HealPixConv(level=6, in_channels=1, out_channels=16,
                   gauge_type="projected_ref",
                   ref_direction=[0, 1, 0])  # singularités à φ=90° et φ=270°

kernelsz=7
# South center
lon=115 
lat=30

level=int(np.log2(nside))
center_idx,_,_=hg.nested.cone_coverage((lon,lat),5.0,level)
acenter_idx,_,_=hg.nested.cone_coverage((lon-180,-lat),5.0,level)

conv = HealPixConv(level=int(np.log2(nside)), in_channels=1, out_channels=1, n_gauges=4, 
                   kernel_sz=kernelsz,gauge_type="projected_ref",
                    singularity_lonlat=(lon,lat))

# 25-point Gaussian: center gets more weight than the 8 neighbours
W = np.array([[0,0,0,0,1,0,0],
              [0,0,0,0,1,1,0],
              [0,0,0,0,1,1,1],
              [1,1,1,1,1,1,1],
              [0,0,0,0,1,1,1],
              [0,0,0,0,1,1,0],
              [0,0,0,0,1,0,0]]).ravel().astype('float32').reshape(1,1,kernelsz*kernelsz)
    
conv.set_kernel(W)

sky = np.zeros([12 * nside**2],dtype=np.float32)
tc,pc=hp.pix2ang(4,np.arange(12*4**2),nest=True)
sky[hp.ang2pix(nside,tc,pc,nest=True)]=1.0
sky_smooth = conv(sky)     # returns np.ndarray [N]
print(sky_smooth.shape)
plt.figure(figsize=(16,10))
for k in range(4):
    tmp=sky_smooth[k]
    tmp[center_idx]=2.
    tmp[acenter_idx]=2.
    mollview(tmp,nest=True,sub=(2,2,1+k),hold=False,coastlines=True,
             title=r'gauge="projected_ref" $\Theta=%d$'%(k))    # (49152,)

(4, 49152)

kernelsz=7
# South center
lon=115 
lat=40

level=int(np.log2(nside))
center_idx,_,_=hg.nested.cone_coverage((-55.0, -10.0),5.0,level)
acenter_idx,_,_=hg.nested.cone_coverage((90.0,  45.0),5.0,level)

center_idx2,_,_=hg.nested.cone_coverage((55.0, 10.0),5.0,level)
acenter_idx2,_,_=hg.nested.cone_coverage((-90.0,  -45.0),5.0,level)

conv = HealPixConv(level=int(np.log2(nside)), in_channels=1, out_channels=1, n_gauges=4, 
                   kernel_sz=kernelsz,
                   gauge_type="two_ref",
                   singularity_lonlat=[
                                (-55.0, -10.0),   # Amazonie  → antipode = (125°E, 10°N) Bornéo
                                ( 90.0,  45.0),   # Afrique   → antipode = (200°E, -10°N) Pacifique central
                                ]
                    )
print(conv.singularity_info())

# 25-point Gaussian: center gets more weight than the 8 neighbours
W = np.array([[0,0,0,0,1,0,0],
              [0,0,0,0,1,1,0],
              [0,0,0,0,1,1,1],
              [1,1,1,1,1,1,1],
              [0,0,0,0,1,1,1],
              [0,0,0,0,1,1,0],
              [0,0,0,0,1,0,0]]).ravel().astype('float32').reshape(1,1,kernelsz*kernelsz)
    
conv.set_kernel(W)

sky = np.zeros([12 * nside**2],dtype=np.float32)
tc,pc=hp.pix2ang(4,np.arange(12*4**2),nest=True)
sky[hp.ang2pix(nside,tc,pc,nest=True)]=1.0
sky_smooth = conv(sky)     # returns np.ndarray [N]
print(sky_smooth.shape)
plt.figure(figsize=(16,10))
for k in range(4):
    tmp=sky_smooth[k]
    tmp[center_idx]=2.
    tmp[acenter_idx]=2.
    tmp[center_idx2]=.5
    tmp[acenter_idx2]=.5
    mollview(tmp,nest=True,sub=(2,2,1+k),hold=False,coastlines=True,
             title=r'gauge="two_ref" $\Theta=%d$'%(k))    # (49152,)

gauge_type='two_ref':
  singularity 1  : lon=+305.00°  lat=-10.00°  (index +1, user-defined)
  singularity 1b : lon=+125.00°  lat=+10.00°  (index +1, antipode of 1)
  singularity 2  : lon=+90.00°  lat=+45.00°  (index +1, user-defined)
  singularity 2b : lon=+270.00°  lat=-45.00°  (index +1, antipode of 2)
  N/S poles      : index -1 each (side-effect, keep outside domain of interest)
  ref_directions : r1=[0.5648625214636235, -0.8067072841115988, -0.17364817766693033]
                   r2=[4.329780281177467e-17, 0.7071067811865476, 0.7071067811865475]
(4, 49152)