Commit 2ecb541b authored by POLCHER Jan's avatar POLCHER Jan 🚴🏾
Browse files

Merge branch 'newdivbas' of gitlab.in2p3.fr:ipsl/lmd/intro/routingpp into newdivbas

parents 1dfeebdf 2ae87fc1
......@@ -34,14 +34,11 @@ function hasanaconda (){
function checkenv (){
rm -f checkenv.py
cat <<EOF > checkenv.py
import mpi4py
import spherical_geometry
from osgeo import ogr
import netCDF4
import matplotlib
import mpl_toolkits
import configparser
import numba
import cartopy
import pyproj
print ("OK")
EOF
python checkenv.py > /dev/null
......@@ -64,42 +61,32 @@ if [ loc != "unknown" ] ; then
#
case $loc in
"ipsl")
module unload python
module load python/3.6-anaconda50
#module unload python
#module load python/3.6-anaconda50
#
# Test if the MPI environment is installed.
#
if [ ! -e ${HOME}/.conda/envs/MPI ] ; then
if [ ! -e ${HOME}/.conda/envs/GEOM ] ; then
#
# Create the MPI environment as it does not exist.
#
conda create -y -n MPI mpi4py numba astropy netcdf4 matplotlib basemap ConfigParser
source activate MPI
conda install -y cartopy
conda install -y --channel astropy spherical-geometry
conda create -y -n GEOM -c conda-forge python=3.7.* cython GDAL pyproj rasterio configparser netcdf4 numba shapely pip
source activate GEOM
pip install argparse
else
source activate MPI
source activate GEOM
fi
;;
"idris")
module load python/3.7.3
if [[ -L ${HOME}/.conda && -d ${HOME}/.conda ]] ; then
if [ ! -e ${WORK}/.conda/envs/MPI ] ; then
conda create -y -n MPI mpi4py numba astropy netcdf4 matplotlib basemap ConfigParser cartopy
conda activate MPI
if [[ -L ${HOME}/.local && -d ${HOME}/.local ]] ; then
export PYTHONUSERBASE=$WORK/.local
if [ $(find $PYTHONUSERBASE -name spherical_geometry -ls | wc -l) -le 0 ] ; then
git clone https://github.com/spacetelescope/spherical_geometry.git
cd spherical_geometry
python setup.py install --user
fi
else
echo "Make sure \$HOME/.link is a link to your \$WORK/.local"
fi
if [ ! -e ${WORK}/.conda/envs/GEOM ] ; then
conda create -y -n GEOM -c conda-forge python=3.7.* cython GDAL pyproj rasterio configparser netcdf4 numba shapely pip
conda activate GEOM
pip install argparse
else
conda activate MPI
fi
conda activate GEOM
fi
else
echo "Make sure \$HOME/.conda is a link to your \$WORK/.conda"
fi
......@@ -112,6 +99,6 @@ if [ loc != "unknown" ] ; then
fi
#
# Verify we can load all the needed modules.
#
#
checkenv
......@@ -3,17 +3,18 @@
#PBS -N remap_MERIT
#
#PBS -j oe
#PBS -l nodes=2:ppn=32
#PBS -l walltime=128:00:00
#PBS -l vmem=300gb,mem=300gb
#PBS -l nodes=1:ppn=42
#PBS -l walltime=06:00:00
#PBS -l vmem=200gb,mem=200gb
#
cd ${PBS_O_WORKDIR}
#
# Set the right Python 3 Anaconda environment
#
source ./Environment
#
source activate GEOM
#source ./Environment_GEOM
\rm *.txt
python ./main_multiprocess.py
python ./MERIT_HydroLakes.py
ls -l
"""
GRID library to handle the different input files
"""
import sys
sys.path.append("../../../")
from netCDF4 import Dataset
import numpy as np
from numba import jit
#from spherical_geometry import polygon
from shapely.geometry import Polygon
import numpy.ma as ma
import time
import sys
import ModelGrid as MG
import RPPtools as RPP
import Projections as PS
import cartopy.io.shapereader as shpreader
from numba import jit
from shapely.geometry import MultiPolygon, Polygon
EarthRadius=6370000.0
import pyproj
from shapely.ops import transform
from osgeo import ogr,gdal
from osgeo import osr
######################################
@jit(nopython = True)
def boxit(cent, dlon, dlat, polyres) :
boxll=[]
loninc = dlon/polyres
latinc = dlat/polyres
# Side 1
for pts in range(polyres) :
boxll.append([cent[0]-dlon/2.0+loninc*pts, cent[1]+dlat/2.0])
# Side 2
for pts in range(polyres) :
boxll.append([cent[0]+dlon/2.0, cent[1]+dlat/2.0-latinc*pts])
# Side 3
for pts in range(polyres) :
boxll.append([cent[0]+dlon/2.0-loninc*pts, cent[1]-dlat/2.0])
# Side 4
for pts in range(polyres) :
boxll.append([cent[0]-dlon/2.0, cent[1]-dlat/2.0+latinc*pts])
# Close
boxll.append(boxll[0])
return boxll
class partition:
def __init__(self, nbcore, nloop, numpolyg):
self.nbcore = nbcore
self.nbloop = nloop
self.numpolyg = numpolyg
n = int(numpolyg / (nbcore*nloop))
self.inter = [[n*i,n*(i+1)] for i in range(nbcore*nloop-1)]
self.inter.append([n*(nbcore*nloop-1),numpolyg])
def get_part(self, i):
return self.inter[i]
class landPolygon:
def __init__(self, land_dir):
self.land_dir = land_dir
self.driver = ogr.GetDriverByName('ESRI Shapefile')
self.dataSource = self.driver.Open(self.land_dir, 0)
self.layer = self.dataSource.GetLayer()
self.nfeature = self.layer.GetFeatureCount()
#
self.Lfeatures = [self.layer.GetFeature(i) for i in range(self.nfeature)]
#####################################
def get_features(self, subLsmall):
self.dataSource = self.driver.Open(self.land_dir, 0)
self.layer = self.dataSource.GetLayer()
self.Lfeatures = [self.layer.GetFeature(i) for i in subLsmall]
class GlobalGrid:
def __init__(self, fdir, input_type):
# Get the coordinates
nc = Dataset(fdir, "r")
#
if input_type == "geogrid":
griddesc, self.lon, self.lat, self.res_lon, self.res_lat = MG.getcoordinates(nc, 0, -1, 0, -1)
self.nj,self.ni = self.lon.shape
# Get the projection
if griddesc['type'] == "lambert_conformal_conic" :
self.proj = PS.LambertC(griddesc['dx'], griddesc['known_lon'], griddesc['known_lat'], griddesc['truelat1'], \
griddesc['truelat2'], griddesc['stdlon'])
elif griddesc['type'] == "rotated_pole" :
self.proj = PS.Cassini(griddesc['lon0'], griddesc['lat0'], griddesc['stdlon'], griddesc['lon1'], griddesc['lat1'], \
griddesc['dlon'], griddesc['dlat'])
elif griddesc['type'] == "RegLonLat" :
self.proj=PS.RegLonLat(self.res_lon, self.res_lat, griddesc['inilon'], griddesc['inilat'])
else :
ERROR("Unknown grid type")
sys.exit()
elif input_type == "hydro":
self.lon=np.copy(nc.variables["nav_lon"][:,:])
self.lat=np.copy(nc.variables["nav_lat"][:,:])
self.nj,self.ni = self.lon.shape
self.res_lon = np.mean(np.abs(np.diff(nc.variables["nav_lon"][0,:])))
self.res_lat = np.mean(np.abs(np.diff(nc.variables["nav_lat"][:,0])))
try:
self.proj=PS.RegLonLat(self.res_lon, self.res_lat, self.lon[0], self.lat[0])
except:
print("Error in Hydrogrid projection")
print("res_lon:{0}, res_lat:{1}, lon0: {2}, lat0: {3}".format(self.res_lon, self.res_lat, self.lon[0], self.lat[0]))
sys.exit()
self.grid_params()
self.get_polygons()
def Unload(self):
for f in self.Lfeatures:
f.Destroy()
del self.Lfeatures
self.dataSource.Destroy()
def grid_params(self):
self.igstart = 0
self.jgstart = 0
class HydroGrid:
def __init__(self, dhydrogrid, extent):
nc = Dataset(dhydrogrid, "r")
self.lon = nc.variables["nav_lon"][0,:]
self.lat = nc.variables["nav_lat"][:,0]
self.hdlon = np.mean(np.abs(np.diff(self.lon)))
self.hdlat = np.mean(np.abs(np.diff(self.lat)))
self.hd = 1* max(abs(self.hdlon), abs(self.hdlat))
# Change name self.ind
indFi=[[i+1] for j in range(self.nj) for i in range(self.ni)]
indFj=[[j+1] for j in range(self.nj) for i in range(self.ni)]
indP_land = [[j,i] for i in range(self.ni) for j in range(self.nj)]
indFi=np.reshape(indFi, self.lon.shape)
indFj=np.reshape(indFj, self.lon.shape)
self.indF = [[indFi[j,i],indFj[j,i]] for j,i in indP_land]
#self.coord = [[self.lon[j,i],self.lat[j,i]] for i,j in self.indF]
#
self.dlon=int((self.lon[0,-1]-self.lon[0,0])/np.abs(self.lon[0,-1]-self.lon[0,0]))
self.dlat=int((self.lat[0,0]-self.lat[-1,0])/np.abs(self.lat[0,0]-self.lat[-1,0]))
# Get the limits of the DOMAIN
i0 = np.argmin(np.abs(self.lon-(extent[0]-self.hd)))
i1 = np.argmin(np.abs(self.lon-(extent[1]+self.hd)))
i0, i1 = np.sort([i0,i1])
j0 = np.argmin(np.abs(self.lat-(extent[2]-self.hd)))
j1 = np.argmin(np.abs(self.lat-(extent[3]+self.hd)))
j0, j1 = np.sort([j0,j1])
self.limits = [j0, j1+1, i0, i1+1]
def get_polygons(self):
print("BORDERS")
t0 = time.time()
Lpolyg = [RPP.boxit(np.array([self.igstart + ij[0],self.jgstart + ij[1] ]), self.dlon, self.dlat, 2) for ij in self.indF]
t1 = time.time()
print("{0} s.".format(int(t1-t0)))
#
print("Corners")
t0 = time.time()
cornersll = [self.proj.ijll(polyg) for polyg in Lpolyg]
t1 = time.time()
print("{0} s.".format(int(t1-t0)))
#
print("Poly")
t0 = time.time()
self.polygg = [Polygon(c) for c in cornersll]
t1 = time.time()
print("{0} s.".format(int(t1-t0)))
self.trip = nc.variables["trip"][j0:j1+1, i0:i1+1]
self.lon = self.lon[i0:i1+1]
self.lat = self.lat[j0:j1+1]
self.ind_land = [[j,i] for j in range(j1+1-j0) for i in range(i1+1-i0)]
######################################
A = np.where(self.trip.mask == False)
if len(A) == 1 :
self.ind_land = [[j,i] for j in range(j1+1-j0) for i in range(i1+1-i0)]
else:
self.ind_land = [[A[0][i],A[1][i]] for i in range(len(A[0]))]
class landPolygon:
def __init__(self, land_dir):
print("Load Land")
t0 = time.time()
shpfi = shpreader.Reader(land_dir)
for k, r in enumerate(shpfi.geometries()):
if k == 0:
self.P = MultiPolygon(r)
else:
self.P.union(MultiPolygon(r))
t1 = time.time()
print("{0} s.".format(int(t1-t0)))
self.get_corners()
nc.close()
#
def get_corners(self):
self.centers = np.array([[self.lon[i], self.lat[j]] for j,i in self.ind_land])
#
def get_poly(self, i):
boxll = boxit(self.centers[i], self.hdlon, self.hdlat, 2)
poly = Polygon(boxll)
if not poly.is_valid:
poly = poly.buffer(0)
return poly
#
def select(self):
indices = np.arange(self.centers.shape[0])
polylist = [self.get_poly(i) for i in indices]
return indices, polylist, self.centers
from GRID import HydroGrid
from intersection import interpolate_polygon, interpolate_polygon_multiproc
def calculate_shapefile(index, landP, dinput, parallel= False):
feature = landP.Lfeatures[index]
shapeid = int(feature.GetField("Hylak_id"))
geom = feature.GetGeometryRef()
extent = geom.GetEnvelope()
del geom; del feature
#
HG = HydroGrid(dinput, extent)
if not parallel:
array = interpolate_polygon(HG,landP,index)
else:
array = interpolate_polygon_multiproc(HG,landP,index)
limits = HG.limits
with open("./out.txt", "a") as foo:
foo.write("part_rank: {0} \n".format(index))
del HG;
return array, limits, shapeid
import numpy as np
import numpy.ma as ma
import time
from osgeo import ogr,osr
from numba import jit
import multiprocessing
from functools import partial
#
EarthRadius=6370000.0
#
......@@ -16,7 +20,6 @@ def interpolate(M, F, option = "conservative"):
# Corresponding MERIT grid points
indices, polylist, subcenters = M.select(c,r)
# F.fp[jfp, ifp] = 1 si sélectionné
OVERLAP = []
if len(indices) > 0:
......@@ -42,5 +45,102 @@ def interpolate(M, F, option = "conservative"):
anc_new[jm,im] = anc_new[jm,im] + frac_modelgrid*value
return anc_new
#
##############################################
def interpolate_polygon(M, landP, ind):
anc_new = ma.zeros(M.trip.shape)
anc_new.mask = M.trip.mask
geom = landP.Lfeatures[ind].GetGeometryRef()
#
indices, polylist, subcenters = M.select()
#
del subcenters;
#
ncell = indices.shape[0]
del indices;
overlap = np.zeros(ncell)
#
chunk = 500
for ll in np.arange(ncell//chunk):
ta = time.time()
t0 = ll*chunk
t1 = t0+chunk
out = [get_overlap(polylist[i],geom) for i in range(t0,t1)]
overlap[t0:t1] = out
tb = time.time()
del out
#
t0 = ncell//chunk
t1 = ncell
overlap[t0:t1] = [get_overlap(polylist[i],geom) for i in range(t0,t1)]
#
A = np.where(overlap > 0)[0][:]
for i in A:
jm, im = M.ind_land[i]
anc_new[jm,im] = anc_new[jm,im] + overlap[i]
return anc_new
######################################
def interpolate_polygon_multiproc(M, landP, ind):
indices, polylist, subcenters = M.select()
#
del subcenters;
#
ncell = indices.shape[0]
del indices;
overlap = np.zeros(ncell)
#
#
chunk = 1000
list_chunk = list(np.arange(ncell//chunk))+[ncell//chunk]
geom = landP.Lfeatures[ind].GetGeometryRef()
wktgeom = geom.ExportToWkt()
del geom
func = partial(local_calculation, chunk = chunk, ncell = ncell, polylist = polylist, wktgeom = wktgeom)
#
nbcore = multiprocessing.cpu_count()
pool = multiprocessing.Pool(processes=nbcore)
P_list = pool.map(func,list_chunk)
pool.close()
#
for out, ll in zip(P_list, list_chunk):
t0 = ll*chunk
t1 = min(t0+chunk, ncell)
overlap[t0:t1] = out
#
#
anc_new = ma.zeros(M.trip.shape)
anc_new.mask = M.trip.mask
A = np.where(overlap > 0)[0][:]
for i in A:
jm, im = M.ind_land[i]
anc_new[jm,im] = anc_new[jm,im] + overlap[i]
return anc_new
####################################
def local_calculation(ll_start, chunk, ncell, polylist, wktgeom):
geom = ogr.CreateGeometryFromWkt(wktgeom)
chunk = 1000
t0 = ll_start*chunk
t1 = min(t0+chunk, ncell)
out = [get_overlap(polylist[i],geom) for i in range(t0,t1)]
return out
def get_overlap(cell, geom):
cellogr = ogr.CreateGeometryFromWkt(cell.wkt)
if cellogr.Intersects(geom):
area_in = cellogr.Intersection(geom).Area()
if area_in > 0:
frac_in = area_in/cellogr.Area()
else:
frac_in = 0
del cellogr;
return frac_in
else:
del cellogr
return 0
......@@ -2,41 +2,66 @@ from netCDF4 import Dataset
import numpy as np
import numpy.ma as ma
def create_output(dinput, P_list,P, gg, output_name):
def create_output(dinput, P_list, output_name, varnamefrac, varnameid):
#
nc = Dataset(dinput, "r")
lon = nc.variables["XLONG_M"][0,:,:]
lat = nc.variables["XLAT_M"][0,:,:]
lon = nc.variables["nav_lon"][0,:]
lat = nc.variables["nav_lat"][:,0]
njg, nig = gg.lon.shape
njg, nig = [lat.shape[0], lon.shape[0]]
# Get the data
outland = np.zeros((njg, nig))
outid = np.zeros((njg, nig))
# Mask the output
# it will be unmasked where we processed the ancillary data
#
# P is partition
for k, out_part in enumerate(P_list):
i0, i1 = P.get_part(k)
for l, m in zip(range(i0,i1),range(len(out_part))):
i,j = gg.indF[l]
outland[j-1,i-1] = out_part[m]
#ancnew = ma.masked_where(ancnew == -1, ancnew)
for out_part, limits,shapeid in P_list:
j0,j1,i0,i1 = limits
b = outid[j0:j1,i0:i1]
b[out_part>0] = shapeid
outid[j0:j1,i0:i1] = b
outland[j0:j1,i0:i1] += out_part
# Start the netCDF
with Dataset(output_name,'w') as foo:
dlat = foo.createDimension('y', (njg))
dlon = foo.createDimension('x', (nig))
#
latitudes = foo.createVariable('lat', np.float32,('y','x'))
longitudes = foo.createVariable('lon', np.float32,('y','x'))
latitudes = foo.createVariable('lat', np.float32,('y'))
longitudes = foo.createVariable('lon', np.float32,('x'))
#
latitudes[:,:] = gg.lat[:,:]
longitudes[:,:] = gg.lon[:,:]
# Only one dimension if regular grid
latitudes[:] = lat[:]
longitudes[:] = lon[:]
#
# Fonction de name_var !!
out = foo.createVariable("landfrac", np.float32,('y','x'))#, zlib = True)
out = foo.createVariable(varnamefrac, np.float32,('y','x'))#, zlib = True)
out[:,:] = outland[:,:]
del outland
out = foo.createVariable(varnameid, np.float32,('y','x'))#, zlib = True)
out[:,:] = outid[:,:]
del outid
foo.sync()
def append_output(dinput, P_list, output_name, varnamefrac,varnameid):
with Dataset(dinput, "r") as nc:
lon = nc.variables["nav_lon"][0,:]
lat = nc.variables["nav_lat"][:,0]
njg, nig = [lat.shape[0], lon.shape[0]]
outland = np.zeros((njg, nig))
outid = np.copy(nc.variables[varnameid][:])
for out_part, limits,shapeid in P_list:
j0,j1,i0,i1 = limits
b = outid[j0:j1,i0:i1]
b[out_part>0] = shapeid
outid[j0:j1,i0:i1] = b
outland[j0:j1,i0:i1] += out_part
with Dataset(output_name,'r+') as foo:
foo.variables[varnamefrac][:,:] += outland
foo.variables[varnameid][:,:] = outid
foo.sync()
del outland,outid
def get_input_shapefileHYDRO(config):
Debug = config.getboolean("OverAll", "Debug", fallback = False)
dinput=config.get("OverAll", "dinput", fallback = None)
shpdir=config.get("OverAll", "shpdir", fallback = None)
input_type=config.get("OverAll", "input_type", fallback = None)
EarthRadius=config.getfloat("OverAll", "EarthRadius", fallback=6370000.0)
#
doutput=config.get("OverAll", "doutput", fallback = None)
varnamefrac=config.get("OverAll", "varnamefrac", fallback = "frac")
varnameid=config.get("OverAll", "varnameid", fallback = "id")
if (input_type is None):
print("Please define input_type")
sys.exit()
elif (input_type not in ["geogrid", "hydro"]):
print("input_type= {0} is not a valid choiche -> = geogrid or = hydro".format(input_type))
elif (dinput is None):
print("Please define dinput")
sys.exit()
elif (shpdir is None):
print("Please define shpdir")
sys.exit()
elif (doutput is None):
print("Please define doutput")
sys.exit()
return Debug,dinput,shpdir,input_type,EarthRadius,doutput,varnamefrac,varnameid
import sys
sys.path.append("./lib")
sys.path.append("../../")
import numpy as np
from netCDF4 import Dataset
import multiprocessing