#* **************************************************************************
#*
#* CDTK, Chemical Dynamics Toolkit
#* A modular system for chemical dynamics applications and more
#*
#* Copyright (C) 2011, 2012, 2013, 2014, 2015, 2016
#* Oriol Vendrell, DESY, <oriol.vendrell@desy.de>
#*
#* Copyright (C) 2017, 2018, 2019
#* Ralph Welsch, DESY, <ralph.welsch@desy.de>
#*
#* Copyright (C) 2020, 2021, 2022, 2023
#* Ludger Inhester, DESY, ludger.inhester@cfel.de
#*
#* This file is part of CDTK.
#*
#* CDTK is free software: you can redistribute it and/or modify
#* it under the terms of the GNU General Public License as published by
#* the Free Software Foundation, either version 3 of the License, or
#* (at your option) any later version.
#*
#* This program is distributed in the hope that it will be useful,
#* but WITHOUT ANY WARRANTY; without even the implied warranty of
#* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
#* GNU General Public License for more details.
#*
#* You should have received a copy of the GNU General Public License
#* along with this program. If not, see <http://www.gnu.org/licenses/>.
#*
#* **************************************************************************
import numpy as np
[docs]
class MultiCounter:
"""Emulate the behaviour of multiple nested do loops
Maintains a list of integers that evolve as the indexes
of a nested "for" loop.
It may be useful to construct nested "for" loops from which
the number of nested elements is not known at programming time,
but at execution time.
"""
def __init__(self,length=1,maxiter=1):
self.pointers=[]
self.ini_pointers=[]
self.limits=[]
self.ini_limits=[]
self.start=True
self.end=False
self.setPointersN(0,length)
self.setLimitsN(maxiter,length)
def __iter__(self):
return self
[docs]
def next(self):
"""Increase the pointers vector and return it
"""
if self.start:
self.start = False
return self.getPointers()
else:
self.increase()
return self.getPointers()
[docs]
def setPointers(self,p):
"""set the list of initial pointers from a given list
"""
self.pointers=[]
self.ini_pointers=[]
for n in p:
self.pointers.append(n)
self.ini_pointers.append(n)
[docs]
def setPointersN(self,n=0,nd=1):
"""set the list of initial pointers to a given number N
n -- number at which the pointers are set
nd -- length of the vector
"""
self.pointers=[]
self.ini_pointers=[]
for i in range(nd):
self.pointers.append(n)
self.ini_pointers.append(n)
[docs]
def setLimits(self,l):
"""set the list of index limits from a given list
"""
self.limits=[]
self.ini_limits=[]
for n in l:
self.limits.append(n)
self.ini_limits.append(n)
[docs]
def setLimitsN(self,n=0,nd=1):
"""set the list of initial limits to a given number N
n -- number at which the pointers are set
nd -- length of the vector
"""
self.limits=[]
for i in range(nd):
self.limits.append(n)
self.ini_limits.append(n)
[docs]
def setMaxIter(self,l):
"""set the list of maximum iterations per index
"""
self.limits=[]
self.ini_limits=[]
for n in l:
self.limits.append(n-1)
self.ini_limits.append(n-1)
[docs]
def setMaxIterN(self,n=1,nd=1):
"""set the list of initial limits to a given number N
n -- number at which the pointers are set
nd -- length of the vector
"""
self.limits=[]
self.ini_limits=[]
for i in range(nd):
self.limits.append(n-1)
self.ini_limits.append(n-1)
[docs]
def getPointers(self):
"""return the list of pointers
"""
return tuple(self.pointers)
[docs]
def getLimits(self):
"""return the list of limits
"""
return tuple(self.limits)
[docs]
def reset(self):
"""reset the iterator to the initial state
"""
self.start=True
self.end=False
for i in range(len(self.pointers)):
self.pointers[i] = self.ini_pointers[i]
self.limits[i] = self.ini_limits[i]
[docs]
def increase(self):
"""increase the list of pointers
return 0 if the end is reached
return 1 if succeed
"""
_increase(self.pointers,self.limits,self.end)
[docs]
class MultiCombination(object):
"""
Iterate over all possible combinations of n indices taken in groups of m
"""
def __init__(self,n,m):
self.n = n
self.m = m
self.mm = m
self.ind = np.zeros((m),int)
self.min = np.zeros((m),int)
self.max = np.zeros((m),int)
for k in xrange(m):
self.ind[k] = k
self.min[k] = k
self.max[k] = n - m + k
if m > 0:
self.ind[-1] = self.ind[-1] - 1
def __iter__(self):
return self
[docs]
def next(self):
if (self.ind==self.max).all():
raise StopIteration
k = self.m - 1
while(self.ind[k] == self.max[k]):
k = k - 1
self.ind[k] += 1
for kk in xrange(k+1,self.m):
self.ind[kk] = self.ind[kk-1] + 1
return self.ind
def _increase(p,l,end):
for i in range(len(p)):
if i == len(p)-1:
if p[i] == l[i]:
end = True
raise StopIteration
if p[i] < l[i]:
p[i] = p[i]+1
break
else:
p[i] = 0
[docs]
def decreaseTuple(t):
"""Decrease in 1 the value at every position of a tuple
"""
l = list(t)
for i in range(len(l)):
l[i]=l[i]-1
return tuple(l)
[docs]
def mapTuple(t1,t2):
"""Return a new tuple with the elements of t1 mapped at the positions
specified in t2
"""
l=[]
for i in t2:
l.append(t1[i])
return tuple(l)
#def excitations(x,m):
# """
# Generate excitation indices for x degrees of freedom up to m quanta
# x -- integer, number of degrees of freedom
# m -- integer, maximum quanta of excitation
# returns an interator object.
# """
# if isinstance(x,int):
# ref = np.ones(x,int)
# yield ref
# else:
# ref = x
# if ref.sum() - len(ref) == m: return
# # make excitations
# exclist = []
# for i in range(len(ref)):
# exc = ref.copy()
# exc[i] = exc[i] + 1
# exclist.append(exc)
# yield exc
# for exc in exclist:
# for exc2 in excitations(exc,m):
# yield exc2
[docs]
def excitations(x,m,n=None,r=None):
"""
Generate excitation indices for x degrees of freedom up to m quanta
x -- integer, number of degrees of freedom
m -- integer, maximum quanta of excitation
----
n -- tuple with maximum states per DOF (max states = max quanta + 1)
r -- reference tuple passed recursively
returns an interator object.
"""
if r is None:
ref = _zeroTuple(x)
yield ref
else:
ref = r
if _sumTuple(ref) == m: return
# make excitations
excdict = {}
for i in range(x):
exc = list(ref)
if n is None:
exc[i] += 1
elif exc[i] < n[i] - 1:
exc[i] += 1
else:
continue
texc = tuple(exc)
excdict[texc] = 0
yield texc
for exc in excdict.keys():
for exc2 in excitations(x,m,n,exc):
if exc2 not in excdict:
excdict[exc2] = 0
yield exc2
[docs]
def excitationsN(x,m,n=None):
"""
Generate excitations exactly of m quanta for x degrees of freedom
x -- integer, number of degrees of freedom
m -- integer, exact quanta of excitation
----
n -- tuple with maximum states per DOF (max states = max quanta + 1)
"""
for e in excitations(x,m,n):
if _sumTuple(e) == m: yield e
def _sumTuple(t):
s = 0
for e in t: s = s + e
return s
def _zeroTuple(l):
z = np.zeros(l,int)
return tuple(z)