In [1]:
from pyspark import SparkContext
ctx = SparkContext('local[*]', 'ccsd')
Or we can also use the dummy spark to emulate the Spark environment in a purely serial way. Note that we need just one Spark context. These two cells should not be both evaluated.
In [ ]:
from dummy_spark import SparkContext
ctx = SparkContext()
With the Spark context, we can construct the drudge specific for this problem. Then we can define some names that is going to be used frequently.
In [2]:
from sympy import *
from drudge import *
dr = PartHoleDrudge(ctx)
dr.full_simplify = False
p = dr.names
c_ = p.c_
c_dag = p.c_dag
a, b = p.V_dumms[:2]
i, j = p.O_dumms[:2]
In [3]:
t = IndexedBase('t')
clusters = dr.einst(
t[a, i] * c_dag[a] * c_[i] +
t[a, b, i, j] * c_dag[a] * c_dag[b] * c_[j] * c_[i] / 4
)
We can have a peek at the cluster operator.
In [4]:
clusters.display()
Out[4]:
Now we need tell the system about the symmetry on $t^2$, so that it can be used in simplification.
In [5]:
dr.set_dbbar_base(t, 2)
Here we can use a loop to nest the commutation conveniently. And IPython magic can be used to time the operation. Note that after the simplification, we explicitly redistribute the terms in the transformed Hamiltonian for better parallel performance in later operations. Note that drudge does not automatically cache the result of tensor computations. The cache method should be called explicitly when a tensor is going to be used multiple times.
In [6]:
%%time
curr = dr.ham
h_bar = dr.ham
for order in range(0, 4):
curr = (curr | clusters).simplify() / (order + 1)
curr.cache()
h_bar += curr
h_bar.repartition(cache=True)
The transformed Hamiltonian can be very complex. Instead of reading its terms, we can just have a peek by get a count of the number of terms it contains.
In [7]:
h_bar.n_terms
Out[7]:
In [8]:
en_eqn = h_bar.eval_fermi_vev().simplify()
We can have a look at its contents to see if it is what we would expect.
In [9]:
en_eqn.display()
Out[9]:
Next, we can create a projector to derive the working equation for the singles amplitude.
In [10]:
proj = c_dag[i] * c_[a]
t1_eqn = (proj * h_bar).eval_fermi_vev().simplify()
In the same way, we can display its content.
In [11]:
t1_eqn.display()
Out[11]:
The working equation for the doubles amplitude can be done in the same way, just it can be slower.
In [12]:
%%time
proj = c_dag[i] * c_dag[j] * c_[b] * c_[a]
t2_eqn = (proj * h_bar).eval_fermi_vev().simplify()
Since the equation can be slightly complex, we can vaguely sort the terms in increasing complexity before display them.
In [13]:
t2_eqn = t2_eqn.sort()
t2_eqn.display()
Out[13]:
Evaluating the working equation takes a lot of effort. Outside drudge, a sister package named gristmill is available for the optimization and automatic code generation for tensor contractions. To start with, we need to put the working equations into a tensor definitions with external indices and import the gristmill package.
In [16]:
from gristmill import *
working_eqn = [
dr.define(Symbol('e'), en_eqn),
dr.define(t[a, i], t1_eqn),
dr.define(t[a, b, i, j], t2_eqn)
]
We can have an estimation of the FLOP cost without any optimization.
In [17]:
orig_cost = get_flop_cost(working_eqn, leading=True)
init_printing()
orig_cost
Out[17]:
Since normally we have far more virtual orbitals than occupied orbitals, we have make the optimization based on this.
In [18]:
%%time
eval_seq = optimize(
working_eqn, substs={p.nv: 5000, p.no: 1000},
contr_strat=ContrStrat.EXHAUST
)
Now we can have some inspection of the evaluation sequence.
In [19]:
len(eval_seq)
Out[19]:
In [20]:
opt_cost = get_flop_cost(eval_seq, leading=True)
opt_cost
Out[20]:
Significant optimization can be seen. Finally we can verify the correctness of the evaluation sequence. This step can be very slow. But it is adviced for mission-critical tasks.
In [21]:
verify_eval_seq(eval_seq, working_eqn, simplify=True)
Out[21]:
Finally, we have have a peek at the details of the intermediates.
In [23]:
for eqn in eval_seq:
eqn.display(False)