mirror of
https://github.com/triqs/dft_tools
synced 2024-12-27 06:43:40 +01:00
203 lines
6.0 KiB
Cython
203 lines
6.0 KiB
Cython
from tools import py_deserialize
|
|
import descriptors
|
|
|
|
cdef class TailGf:
|
|
cdef tail _c
|
|
def __init__(self, **d):
|
|
"""
|
|
TailGf ( shape )
|
|
TailGf ( data, mask, order_min )
|
|
"""
|
|
c_obj = d.pop('encapsulated_c_object', None)
|
|
if c_obj :
|
|
assert d == {}
|
|
self._c = extractor [tail] (c_obj) ()
|
|
return
|
|
|
|
bss = d.pop('boost_serialization_string', None)
|
|
if bss :
|
|
assert d == {}, "Internal error : boost_serialization_string must be the only argument"
|
|
boost_unserialize_into(<std_string>bss,self._c)
|
|
return
|
|
|
|
# default values
|
|
omin = -1
|
|
omax = 8
|
|
|
|
a = d.pop('data',None)
|
|
if a==None :
|
|
(N1, N2) = d.pop('shape')
|
|
a = numpy.zeros((omax-omin+1,N1,N2), numpy.complex)
|
|
m = d.pop('mask',None)
|
|
if m==None :
|
|
m = numpy.zeros(a.shape[1:3], int)
|
|
m.fill(omax)
|
|
o = d.pop('order_min',None)
|
|
if o==None: o = omin
|
|
|
|
assert len(d) == 0, "Unknown parameters in TailGf constructions %s"%d.keys()
|
|
|
|
self._c = tail(array_view[dcomplex,THREE](a), array_view[long,TWO](m), o)
|
|
|
|
#-------------- Reduction -------------------------------
|
|
|
|
def __reduce__(self):
|
|
return py_deserialize, (self.__class__,boost_serialize(self._c),)
|
|
|
|
#-------------- Properties -------------------------------
|
|
|
|
property data :
|
|
"""Access to the data array"""
|
|
def __get__(self) :
|
|
return self._c.data().to_python()
|
|
|
|
property mask:
|
|
"""Access to the mask"""
|
|
def __get__(self) :
|
|
return self._c.mask_view().to_python()
|
|
|
|
property shape :
|
|
def __get__(self) : return self.data.shape[1:3]
|
|
|
|
property order_min :
|
|
"""Min order of the expansion"""
|
|
def __get__(self) : return self._c.order_min()
|
|
|
|
property order_max :
|
|
"""Max order of the expansion"""
|
|
def __get__(self) : return self._c.order_max()
|
|
|
|
property N1 :
|
|
def __get__(self): return self.shape[0]
|
|
|
|
property N2 :
|
|
def __get__(self): return self.shape[1]
|
|
|
|
property size :
|
|
"""Length of the expansion"""
|
|
def __get__(self) : return self._c.size()
|
|
|
|
def copy(self) :
|
|
return self.__class__(data = self.data.copy(), mask = self.mask.copy())
|
|
|
|
def copy_from(self, TailGf T) :
|
|
self._c << T._c
|
|
|
|
def _make_slice(self, sl1, sl2):
|
|
return self.__class__(data = self.data[:,sl1,sl2], mask = self.mask[sl1,sl2])
|
|
|
|
def __repr__ (self) :
|
|
omin = self.order_min
|
|
while ((omin <= self.order_max) and (numpy.max(numpy.abs(self.data[omin-self.order_min,:,:])) < 1e-8)):
|
|
omin = omin+1
|
|
if omin == self.order_max+1:
|
|
return "%s"%numpy.zeros(self.shape)
|
|
else:
|
|
return string.join([ "%s"%self[r]+ (" /" if r>0 else "") + " Om^%s"%(abs(r)) for r in range(omin, self.order_max+1) ] , " + ")
|
|
|
|
def __getitem__(self,i) :
|
|
"""Returns the i-th coefficient of the expansion, or order Om^i"""
|
|
if not self.has_coef(i) : raise IndexError, "Index %s is out of range"%i
|
|
return self.data[i-self.order_min,:,:]
|
|
|
|
def __setitem__(self,i, val) :
|
|
"""Sets the i-th coefficient of the expansion, or order Om^i"""
|
|
if not self.has_coef(i) : raise IndexError, "Index %s is out of range"%i
|
|
self.data[i-self.order_min,:,:] = val
|
|
|
|
def has_coef(self, i):
|
|
return (i >= self.order_min) and (i <= self.order_max)
|
|
|
|
def __call__(self, x) :
|
|
val = 0.0
|
|
for n in range(self.order_min, self.order_max+1):
|
|
val += self[n] * x**(-n)
|
|
return val
|
|
|
|
def invert(self) :
|
|
self._c << inverse_c (self._c)
|
|
|
|
######### arithmetic operations #################
|
|
|
|
def __iadd__(self, TailGf arg):
|
|
self._c << self._c + arg._c
|
|
return self
|
|
|
|
def __add__(self, TailGf y):
|
|
cdef tail t = (<TailGf?>self)._c + y._c
|
|
return make_TailGf(t)
|
|
#c = self.copy()
|
|
#c += y
|
|
#return c
|
|
|
|
def __isub__(self, TailGf arg):
|
|
self._c << self._c - arg._c
|
|
return self
|
|
|
|
def __sub__(self,TailGf y):
|
|
cdef tail t = (<TailGf?>self)._c - y._c
|
|
return make_TailGf(t)
|
|
#c = self.copy()
|
|
#c -= y
|
|
#return c
|
|
|
|
def __imul__(self,arg):
|
|
n = type(arg).__name__
|
|
if n == 'TailGf' :
|
|
self._c << self._c * (<TailGf?>arg)._c
|
|
elif descriptors.is_scalar(arg):
|
|
self._c << as_dcomplex(arg)* self._c
|
|
else :
|
|
raise RuntimeError, " argument type not recognized in imul for %s"%arg
|
|
return self
|
|
|
|
def __mul__(self,arg):
|
|
c = self.copy()
|
|
c *= arg
|
|
return c
|
|
|
|
def __idiv__(self,arg):
|
|
if descriptors.is_scalar(arg):
|
|
self._c << self._c / as_dcomplex(arg)
|
|
else:
|
|
raise RuntimeError, "rhs must be a scalar"
|
|
return self
|
|
|
|
def __div__(self,arg):
|
|
c = self.copy()
|
|
c /= arg
|
|
return c
|
|
|
|
#---- other operations ----
|
|
|
|
def zero(self) :
|
|
"""Sets the expansion to 0"""
|
|
self._c << as_dcomplex(0.0)
|
|
|
|
def transpose (self) :
|
|
"""Transpose the array : new view as in numpy"""
|
|
return TailGf(data=self.data.transpose(), mask=self.mask.transpose())
|
|
|
|
def conjugate(self) :
|
|
"""Transpose the array : new view as in numpy"""
|
|
return TailGf(data=self.data.conjugate(), mask=self.mask)
|
|
|
|
def __write_hdf5__ (self, gr , char * key) :
|
|
h5_write (make_h5_group(gr), key, self._c)
|
|
|
|
#---------------- Reading from h5 ---------------------------------------
|
|
|
|
def h5_read_TailGf( gr, std_string key) :
|
|
return make_TailGf( h5_extractor[tail]()(make_h5_group(gr),key))
|
|
|
|
from pytriqs.archive.hdf_archive_schemes import register_class
|
|
register_class (TailGf, read_fun = h5_read_TailGf)
|
|
|
|
#-----------------------------------------------------
|
|
# C -> Python
|
|
#-----------------------------------------------------
|
|
|
|
cdef inline make_TailGf ( tail x) :
|
|
return TailGf(encapsulated_c_object = encapsulate (&x))
|
|
|