Changeset 640

Timestamp:

12/18/17 12:43:01 (6 years ago)

Author:

dubos

Message:

devel/unstructured : Python interface to ARK time scheme

File:

• codes/icosagcm/devel/Python/src/unstructured.pyx (modified) (3 diffs)

codes/icosagcm/devel/Python/src/unstructured.pyx

@@ -5 +5 @@
 from ctypes import c_void_p, c_int, c_double, c_bool
 
+from libc.time cimport time as ctime, time_t
+cimport numpy as np
+
 #------------- direct Cython interface to DYNAMICO routines -------------#
 
+cdef enum: max_nb_stage=5
+cdef extern :
+    cdef double tauj[max_nb_stage]
+    cdef double cslj[max_nb_stage][max_nb_stage]
+    cdef double cflj[max_nb_stage][max_nb_stage]
+    cdef int nb_stage[1]
+
 cdef extern from "functions.h":
+     cdef void dynamico_ARK_step(double *mass_col, double *rhodz, double *theta_rhodz, 
+                                 double *u, double *geopot, double *w,
+                                 double *theta, double *ps, double *pk, double *hflux, double *qv,
+                                 double *dmass_col, double *drhodz, double *dtheta_rhodz,
+                                 double *du_fast, double *du_slow,
+                                 double *dPhi_fast, double *dPhi_slow, 
+                                 double *dW_fast, double *dW_slow)
      cdef void dynamico_init_params()
      cpdef void dynamico_setup_xios()
@@ -46 +63 @@
 
 kernels.addvars(
-    c_bool,'hydrostatic','debug_hevi_solver','rigid',
+    c_bool,'hydrostatic','debug_hevi_solver',
     c_int,'llm','nqdyn','primal_num','max_primal_deg',
     'dual_num','max_dual_deg','edge_num','max_trisk_deg',
@@ -54 +71 @@
 
 elapsed=0.
+
+#------------------------ Extension type performing a full ARK time step ----------------------
+
+ctypedef double *dbl_ptr
+cdef dbl_ptr ptr1(double[:] data) : return &data[0]
+cdef dbl_ptr ptr2(double[:,:] data) : return &data[0,0]
+cdef dbl_ptr ptr3(double[:,:,:] data) : return &data[0,0,0]
+cdef dbl_ptr ptr(data):
+    n=data.ndim
+    if n==1 : return ptr1(data)
+    if n==2 : return ptr2(data)
+    if n==3 : return ptr3(data)
+
+cdef alloc(dbl_ptr *p, allocator, n=1):
+    data=allocator(n)
+    p[0]=ptr(data)
+    return data
+
+cdef check_ptr(name, dbl_ptr p, np.ndarray data):
+    if p != ptr(data) : print name, 'p <> ptr(data) !!'
+        
+cdef class Caldyn_step:
+    # pointer to allocated arrays
+    cdef dbl_ptr p_mass, p_theta_rhodz, p_u, p_geopot, p_w                   # prognostic
+    cdef dbl_ptr p_mass_col, p_dmass_col, p_ps, p_theta, p_pk, p_hflux, p_qv # diagnostic
+    cdef dbl_ptr p_drhodz, p_dtheta_rhodz, p_du_fast, p_du_slow                # tendencies
+    cdef dbl_ptr p_dPhi_fast, p_dPhi_slow, p_dW_fast, p_dW_slow                # tendencies
+    # allocated arrays, must remain referenced or segfault
+    cdef readonly np.ndarray mass, theta_rhodz, u, geopot, w
+    cdef readonly np.ndarray mass_col, dmass_col, ps, theta, pk, hflux, qv
+    cdef readonly np.ndarray drhodz, dtheta_rhodz, du_fast, du_slow
+    cdef readonly np.ndarray dPhi_fast, dPhi_slow, dW_fast, dW_slow
+
+    def __init__(self,mesh,time_scheme):
+        #        self.mesh=mesh
+        fps, ftheta, fmass = mesh.field_ps, mesh.field_theta, mesh.field_mass
+        fw, fu, fz = mesh.field_w, mesh.field_u, mesh.field_z
+        # collect coefficients of time scheme
+        cdef double[:]   tauj_ = time_scheme.tauj
+        cdef double[:,:] cslj_ = time_scheme.csjl
+        cdef double[:,:] cflj_ = time_scheme.cfjl
+        ns = time_scheme.nstage
+        nb_stage[0]= ns
+        cdef int i,j
+        for i in range(ns):
+            tauj[i]=tauj_[i]
+            for j in range(ns):
+                cslj[i][j]=cslj_[i,j]
+                cflj[i][j]=cflj_[i,j]
+        # allocate arrays, store pointers to avoid overhead when calling dynamico
+        #    prognostic/diagnostic
+        self.ps                       = alloc(&self.p_ps, fps) 
+        self.mass_col, self.dmass_col = alloc(&self.p_mass_col, fps), alloc(&self.p_dmass_col, fps), 
+        self.mass, self.theta_rhodz   = alloc(&self.p_mass, fmass), alloc(&self.p_theta_rhodz, fmass),
+        self.theta, self.pk           = alloc(&self.p_theta, fmass), alloc(&self.p_pk, fmass), 
+        self.geopot, self.w           = alloc(&self.p_geopot, fw), alloc(&self.p_w, fw),
+        self.hflux, self.u            = alloc(&self.p_hflux, fu), alloc(&self.p_u, fu) 
+        self.qv                       = alloc(&self.p_qv,fz)
+        #    tendencies
+        self.drhodz, self.dtheta_rhodz  = alloc(&self.p_drhodz,fmass,ns), alloc(&self.p_dtheta_rhodz,fmass,ns) 
+        self.du_fast, self.du_slow      = alloc(&self.p_du_fast,fu,ns), alloc(&self.p_du_slow,fu,ns)
+        self.dPhi_fast, self.dPhi_slow  = alloc(&self.p_dPhi_fast,fw,ns), alloc(&self.p_dPhi_slow,fw,ns)
+        self.dW_fast, self.dW_slow      = alloc(&self.p_dW_fast,fw,ns), alloc(&self.p_dW_slow,fw,ns)
+    def next(self):
+        #        global elapsed
+        # time1=time.time()
+        check_ptr('mass', self.p_mass, self.mass)
+        check_ptr('u', self.p_u, self.u)
+        dynamico_ARK_step(self.p_mass_col, self.p_mass, self.p_theta_rhodz,
+                          self.p_u, self.p_geopot, self.p_w,
+                          self.p_theta, self.p_ps, self.p_pk, self.p_hflux, self.p_qv,
+                          self.p_dmass_col, self.p_drhodz, self.p_dtheta_rhodz,
+                          self.p_du_fast, self.p_du_slow,
+                          self.p_dPhi_fast, self.p_dPhi_slow,
+                          self.p_dW_fast, self.p_dW_slow)
+        #time2=time.time()
+        #if time2>time1: elapsed=elapsed+time2-time1
+    def setup_TRSW(self):
+        setvars(('hydrostatic','caldyn_thermo','caldyn_eta'),
+                (True,thermo_boussinesq,eta_lag))
+        dynamico_init_params()
+    def setup_HPE(self, caldyn_thermo, caldyn_eta, g, BC,thermo):
+        setvars(('hydrostatic','caldyn_thermo','caldyn_eta',
+                 'g','ptop','Rd','cpp','preff','Treff'),
+                (True,caldyn_thermo,caldyn_eta,
+                 g,BC.ptop,thermo.Rd,thermo.Cpd,thermo.p0,thermo.T0))
+        dynamico_init_params()
+    def setup_NH(self, caldyn_thermo, caldyn_eta, g, BC,thermo):
+        setvars(('hydrostatic','caldyn_thermo','caldyn_eta',
+                 'g','ptop','Rd','cpp','preff','Treff',
+                 'pbot','rho_bot'),
+                (False,caldyn_thermo,caldyn_eta,
+                 g,BC.ptop,thermo.Rd,thermo.Cpd,thermo.p0,thermo.T0,
+                 BC.pbot.max(), BC.rho_bot.max()))
+        dynamico_init_params()
 
 #----------------------------- Base class for dynamics ------------------------