fix faulty polygon tests for grid with 180 meridians for shtns backend.

Author: strawpants

src/shxarray/geom/polygons.py

@@ -41,6 +41,7 @@ def polygon2sh(polygeom,nmax:int=100,auxcoord=None,engine="shlib",**kwargs) ->xr
     if type(polygeom) != gpd.GeoSeries:
         polygeom=gpd.GeoSeries(polygeom)
 
+    
     #create a dense enough grid encompassing all polgyons to use for spherical harmonic synthesis
     # heuristic way to figure out the resolution based on nmax
     dslonlat=xr.Dataset.sh.lonlat_grid(nmax,engine=engine)
@@ -66,9 +67,14 @@ def polygon2sh(polygeom,nmax:int=100,auxcoord=None,engine="shlib",**kwargs) ->xr
 
 
     dtmp=xr.DataArray(np.zeros([dslonlat.sizes['lon'],dslonlat.sizes['lat'],len(polygeom)]),coords=coords,dims=dims).stack(lonlat=("lon","lat"))
-
-    #create a geoDataframe of points from the grid
-    ggrd=gpd.GeoDataFrame(geometry=[Point(lon,lat) for lon,lat in dtmp.lonlat.values],crs=4326)
+    if dtmp.lon.min() < 0:
+        #grid has 0 central meridian already
+        #create a geoDataframe of points from the grid (lon is already with 0 central meridian
+        ggrd=gpd.GeoDataFrame(geometry=[Point(lon,lat) for lon,lat in dtmp.lonlat.values],crs=4326)
+    else:
+        #grid has 180 central meridian
+        #create a geoDataframe of points from the grid and convert lon to have 0 central meridian
+        ggrd=gpd.GeoDataFrame(geometry=[Point((lon+180)%360-180,lat) for lon,lat in dtmp.lonlat.values],crs=4326)
 
     if polygeom.crs != ggrd.crs:
         #possibly convert the lon/lat grid in the desired projection before doing the polygon test
@@ -77,7 +83,11 @@ def polygon2sh(polygeom,nmax:int=100,auxcoord=None,engine="shlib",**kwargs) ->xr
 
     #query using a spatial index and set values to 1
     shxlogger.info("Masking and gridding polygons")
-    for i,poly in enumerate(polygeom): 
+    for i,poly in enumerate(polygeom):
+        # if i == 18:
+            # breakpoint()
+            # from IPython.core.debugger import set_trace
+            # set_trace()
         idx=ggrd.sindex.query(poly,predicate="contains")
         dtmp[i,idx]=1.0
 

src/shxarray/kernels/anisokernel.py

@@ -8,7 +8,7 @@
 from shxarray.shlib import Ynm
 from shxarray.core.sh_indexing import SHindexBase
 import sparse
-
+import numpy as np
 from packaging import version
 
 
@@ -48,8 +48,12 @@ def __call__(self,dain:xr.DataArray):
 
         if self.nmax < dain.sh.nmax:
             raise RuntimeError("Input data has higher degree than kernel, cannot apply kernel operator to object")
+        
+        if type(dain.data) != np.ndarray:
+            dain=dain.compute()
 
         daout=xr.dot(self._dskernel.mat,dain,dims=[SHindexBase.name]) 
+
         #rename nm and convert to dense array
         daout=daout.sh.toggle_nm()
         if self.useDask:

src/shxarray/signal/basinav.py

@@ -44,5 +44,5 @@ def __call__(self, datws,**kwargs):
             leakage=leakage_corr_vishwa2016(datws, self._dabin, self._filtername,engine=engine) 
             da_av=(da_av-leakage)*dascales
 
-        return da_av
+        return da_av.drop_vars(['n','m','nm'])
 

src/shxarray/signal/leakage_vishwa2016.py

@@ -8,7 +8,7 @@
 from scipy.signal import hilbert
 from shxarray.kernels import getSHfilter
 from shxarray.exp.multiply import multiply
-
+from shxarray.core.logging import shxlogger
 
 def leakage_corr_vishwa2016(datws, dabasins, filtername,engine='shlib'):
     """
@@ -59,9 +59,8 @@ def leakage_corr_vishwa2016(datws, dabasins, filtername,engine='shlib'):
     filterOp = getSHfilter(filtername,nmax=datws.sh.nmax)
     datws_f=filterOp(datws)
 
-    
+     
     Ic_f=((daknm@datws_f)/dabasins.sel(n=0,m=0)).interp({timedim:time_eq})
-
     #compute double filtered leakage signal
     datws_ff=filterOp(datws_f)
     Ic_ff=((daknm@datws_ff)/dabasins.sel(n=0,m=0)).interp({timedim:time_eq})
@@ -79,7 +78,11 @@ def leakage_corr_vishwa2016(datws, dabasins, filtername,engine='shlib'):
         A[:,0]=1
         A[:,1]=Ic_ff.isel({auxdim:i}).data
         A[:,2]=np.imag(hilbert(Ic_ff.isel({auxdim:i}).data))
-        reg_fit,_,_,_=np.linalg.lstsq(A,Ic_f.isel({auxdim:i}).data)
+        try:
+            reg_fit,_,_,_=np.linalg.lstsq(A,Ic_f.isel({auxdim:i}).data)
+        except np.linalg.LinAlgError:
+            shxlogger.warning(f"Least squares fit failed for basin {i}, phase for leakage could not be computed, assuming np phase shift")
+            reg_fit=[0,1,0]
         # compute phase=atan(c/b) to get the phase shift and take the complex exponential
         phase_exp.append(np.exp(-1j*(np.atan2(reg_fit[2],reg_fit[1]))))