text:
import numpy as np
import pandas as pd
from tspdb.src.database_module.db_class import Interface
from tspdb.src.database_module.plpy_imp import plpyimp
#######################TO DO##########################
#1 get SUV instead of all getU,getS, getV
######################################################
class plpyimp(plpyimp):

    def get_time_series(self, name, start, end = None, start_ts = '1970/01/01 00:00:00', value_column="ts", index_column='row_id', Desc=False, interval = 60, aggregation_method = 'average' ):

        """
        query time series table to return equally-spaced time series values from a certain range  [start to end]
        or all values with time stamp/index greater than start  (if end is None)
        ----------
        Parameters
        ----------
        name: string
            table (time series) name in database

        start: int or  timestamp
            start index (timestamp) of the range query

        end: int, timestamp
            last index (timestamp) of the range query

        value_column: string
            name of column than contain time series value

        index_column: string
            name of column that contains time series index/timestamp

        interval: float optional (default=60)
            if time-index type is timestamp, determine the period (in seconds) in which the timestamps are truncated

        aggregation_method: str optional (default='average')
            the method used to aggragte values belonging to the same interval. options are: 'average', 'max', 'min', and 'median'

        desc: boolean optional (default=false)
            if true(false),  the returned values are sorted descendingly (ascendingly) according to index_column
        ----------
        Returns
        ----------
        array, shape [(end - start +1) or  ceil(end(in seconds) - start(in seconds) +1) / interval ]
            Values of time series in the time interval start to end sorted according to index_col
        """
        # check if hypertable

        hypertable = self.engine.execute("SELECT count(*)>=1 as h FROM timescaledb_information.hypertable WHERE table_schema='public' AND table_name='%s';" % name)[0]['h']

        if isinstance(start, (int, np.integer)) and (isinstance(end, (int, np.integer)) or end is None):
            if end is None:
                sql = 'Select ' + value_column + " from  " + name + " where " + index_column + " >= "+str(start)+" order by "+index_column
                result = self.engine.execute(sql)

            else:
                if not Desc:
                    sql = 'Select ' + value_column + " from  " + name + " where " + index_column + " >= "+str(start)+" and " + index_column + " <= "+str(end)+" order by " + index_col
                else:
                    sql = 'Select ' + value_column + " from  " + name + " where " + index_column + " >= "+str(start)+" and " + index_column + " <= "+str(end)+" order by " + index_col + ' Desc'
                result = self.engine.execute(sql)
            result = [row for row in result]
            columns = value_column.split(',')
            return [[row[ci] for ci in columns] for row in result]

        elif  isinstance(start, (pd.Timestamp)) and (isinstance(end, (pd.Timestamp)) or end is None):
            #SELECT
            #time_bucket_gapfill('00:00:05', time) AS date,
            #avg(ts)
            #FROM ts_basic_ts_5_5
            #WHERE time >= '1/10/2012' AND time < ' 2012-10-06 18:53:15'
            #GROUP BY date
            #ORDER BY date;

            seconds = interval%60
            minutes = int(interval/60)
            hours = int(interval/3600)
            interval_str = '%s:%s:%s'%(hours, minutes, seconds)

            agg_func_dict = {'average': 'AVG', 'min': 'MIN', 'max': 'MAX'}
            try:
                agg_function = agg_func_dict[aggregation_method]
            except KeyError as e:
                print ('aggregation_method not valid choose from ("average", "min", "max"), Exception: "%s"' % str(e))
                raise
            ## might be needed
            start_ts_str = start_ts.strftime('%Y-%m-%d %H:%M:%S')
            ## fix strings formatting
            if hypertable:
                if end is None:
                    sql = "SELECT time_bucket_gapfill('%s', "+index_column+") AS date, "+agg_function+"("+value_column+") as avg_val FROM "+name+" where "+index_column+" >= '%s' and  "+index_column+" <= '%s' GROUP BY date ORDER BY date"
                    sql = sql%(interval_str, start.strftime('%Y-%m-%d %H:%M:%S'), 'now()',)
                else:
                    sql = "SELECT time_bucket_gapfill('%s', "+index_column+") AS date, "+agg_function+"("+value_column+") as avg_val FROM "+name+" where "+index_column+" >= '%s' and  "+index_column+" <= '%s' GROUP BY date ORDER BY date"
                    sql = sql%(interval_str, start.strftime('%Y-%m-%d %H:%M:%S'), end.strftime('%Y-%m-%d %H:%M:%S'),)
            else:
                if end is None:
                    select_sql = "select "+agg_function+"(m."+value_column+") avg_val from "+name+" m right join intervals f on m."+index_column+" >= f.start_time and m."+index_column+" < f.end_time where f.end_time > "+start_ts_str+"  group by f.start_time, f.end_time order by f.start_time"
                    generate_series_sql = "with intervals as (select n as start_time,n+'"+interval_str+"'::interval as end_time from generate_series('"+start.strftime('%Y-%m-%d %H:%M:%S')+"'::timestamp, now(),'"+interval_str+"'::interval) as n )"
                else:
                    generate_series_sql = "with intervals as (select n as start_time,n+'"+interval_str+"'::interval as end_time from generate_series('%s'::timestamp, '%s'::timestamp,'"+interval_str+"'::interval) as n )"
                    generate_series_sql = generate_series_sql % (start_ts_str,end.strftime('%Y-%m-%d %H:%M:%S'),)
                    select_sql = "select "+agg_function+"(m."+value_column+") avg_val from "+name+" m right join intervals f on m."+index_column+" >= f.start_time and m."+index_column+" < f.end_time where f.end_time > '%s' and  f.start_time < '%s' group by f.start_time, f.end_time order by f.start_time"
                    select_sql = select_sql%(start.strftime('%Y-%m-%d %H:%M:%S'), end.strftime('%Y-%m-%d %H:%M:%S'),)

                sql = generate_series_sql+ select_sql
            if Desc: sql += 'DESC'
            result = self.engine.execute(sql)
            result = [row for row in result]
            return [(row['avg_val'],) for row in result]

        else:
             raise Exception('start and end values must either be integers or pd.timestamp')


text:
import numpy as np
import pandas as pd

from quantopian.pipeline import Pipeline
from quantopian.pipeline.domain import US_EQUITIES
from quantopian.pipeline.data import EquityPricing
from quantopian.pipeline.data.factset import Fundamentals, EquityMetadata, RBICSFocus
from quantopian.pipeline.factors import (
    CustomFactor,
    AnnualizedVolatility,
    SimpleBeta,
    Returns,
    SimpleMovingAverage,
    AverageDollarVolume,
)

# import groupings
# ----------------
rbics_economy_number = RBICSFocus.l1_id.latest
rbics_economy = RBICSFocus.l1_name.latest

# Defining universe criteria
# --------------------------
avg_day_dv_200 = AverageDollarVolume(window_length=200)
mcap = Fundamentals.mkt_val.latest
price = EquityPricing.close.latest
volume = EquityPricing.volume.latest

# Set Universe
# ------------
universe = (
    avg_day_dv_200.percentile_between(5, 100)
    & (price > 5.0)
    & (mcap > 100e6)
    & EquityMetadata.security_type.latest.eq("SHARE")
    & EquityMetadata.is_primary.latest
    & volume.notnull()
)

universe_top500 = mcap.top(500, mask=universe)

# Creating custom factors
# -----------------------
class lt_mom(CustomFactor):
    inputs = [EquityPricing.close]
    window_length = 252

    def compute(self, today, assets, out, price):
        out[:] = (price[-21] - price[-252]) / price[-252]

# Size
# ----
mcap = Fundamentals.mkt_val.latest
log_mcap = (
    mcap.log()
    .winsorize(min_percentile=0.05, max_percentile=0.95, mask=universe_top500)
    .zscore()
)

# Value
# -----
fcfy = Fundamentals.free_cf_fcfe_ltm.latest / mcap
fcfy = fcfy.winsorize(
    min_percentile=0.05, max_percentile=0.95, mask=universe_top500
).zscore()

ey = Fundamentals.oper_inc_bef_dep_ltm.latest / mcap
ey = ey.winsorize(
    min_percentile=0.05, max_percentile=0.95, mask=universe_top500
).zscore()

# Momentum
# --------
lt_momentum = lt_mom(mask=universe_top500).zscore()

# Volatility
# ----------
ann_vol = AnnualizedVolatility(
    inputs=[Returns(window_length=2)], window_length=66, mask=universe_top500
).zscore()

# Quality
# -------
gpa = Fundamentals.gross_inc_ltm.latest / Fundamentals.assets.latest
gpa = gpa.zscore(mask=universe_top500)

factors = {
    "fcfy": fcfy,
    "ey": ey,
    "log_mcap": log_mcap,
    "lt_mom": lt_momentum,
    "ann_vol": ann_vol,
    "gpa": gpa,
}

def make_pipeline():

    return Pipeline(
        columns=factors,
        screen=universe_top500.downsample("month_start"),
        domain=US_EQUITIES,
    )

ext: py

max_stars_repo_name: quantopian/aqueduct-client

max_issues_repo_name: quantopian/aqueduct-client

max_forks_repo_name: quantopian/aqueduct-client

reason: import numpy


text:
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_cycles:
#
# Color cycles
# ============
#
# Proplot defines **color cycles** or **discrete colormaps** as color palettes
# comprising sets of *distinct colors*. Unlike :ref:`continuous colormaps <ug_cmaps>`,
# interpolation between these colors may not make sense. Generally, color cycles are
# used with line plots, bar plots, and other distinct plot elements. Occasionally,
# they are used as colormaps for qualitative or categorical data. Proplot's
# color cycles are registered as `~proplot.colors.DiscreteColormap`\ s,
# and can be converted into matplotlib `property cyclers
# <https://matplotlib.org/stable/tutorials/intermediate/color_cycle.html>`__
# for use with distinct plot elements using the `~proplot.constructor.Cycle`
# constructor function. `~proplot.constructor.Cycle` can also
# :ref:`extract colors <ug_cycles_new>` from `~proplot.colors.ContinuousColormap`\ s.
#
# Proplot :ref:`adds several features <why_colormaps_cycles>` to help you use color
# cycles effectively in your figures. This section documents the new registered
# color cycles, explains how to make and modify color cycles, and shows how to
# apply them to your plots.

# %% [raw] raw_mimetype="text/restructuredtext" tags=[]
# .. _ug_cycles_included:
#
# Included color cycles
# ---------------------
#
# Use `~proplot.demos.show_cycles` to generate a table of the color cycles
# registered by default, loaded from the ``cycles``
# `~proplot.config.Configurator.user_folder`, and/or created with the
# `~proplot.constructor.Cycle` constructor function. To retrieve the list
# of colors associated with a registered or on-the-fly color cycle,
# simply use `~proplot.utils.get_colors`.

# %%
import proplot as pplt
fig, axs = pplt.show_cycles()

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_cycles_changing:
#
# Changing the color cycle
# ------------------------
#
# Various plotting commands like `~proplot.axes.PlotAxes.line` and
# `~proplot.axes.PlotAxes.scatter` now accept a `cycle` keyword
# passed to the `~proplot.constructor.Cycle` constructor function
# (see the :ref:`1D plotting section <ug_apply_cycle>`). To save
# your color cycle data and use it every time proplot is imported, simply pass
# ``save=True`` to `~proplot.constructor.Cycle`. If you want to change the global
# property cycler, pass a `~proplot.colors.DiscreteColormap` or colormap name
# to :rcraw:`cycle` or pass the result of `~proplot.constructor.Cycle`
# to :rcraw:`axes.prop_cycle` (see the :ref:`configuration guide <ug_config>`).

# %%
import proplot as pplt
import numpy as np

# Sample data
state = np.random.RandomState(51423)
data = (state.rand(12, 6) - 0.45).cumsum(axis=0)
kwargs = {'legend': 'b', 'labels': list('abcdef')}

# Figure
lw = 5
pplt.rc.cycle = '538'
fig = pplt.figure(refwidth=1.9)
fig.format(suptitle='Changing the color cycle')

# Modify the default color cycle
ax = fig.subplot(131)
ax.plot(data, lw=lw, **kwargs)
ax.format(title='Global color cycle')

# Pass the cycle to a plotting command
ax = fig.subplot(132)
ax.plot(data, cycle='qual1', lw=lw, **kwargs)
ax.format(title='Local color cycle')

# As above but draw each line individually
# Note that passing cycle=name to successive plot calls does
# not reset the cycle position if the cycle is unchanged
ax = fig.subplot(133)
labels = kwargs['labels']
for i in range(data.shape[1]):
    ax.plot(data[:, i], cycle='qual1', legend='b', label=labels[i], lw=lw)
ax.format(title='Multiple plot calls')

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_cycles_new:
#
# Making color cycles
# -------------------
#
# You can make new color cycles with the `~proplot.constructor.Cycle`
# :ref:`constructor function <why_constructor>`. One great way to make cycles is by
# sampling colormaps! Just pass the colormap name to `~proplot.constructor.Cycle`,
# and optionally specify the number of samples you want to draw as the last
# positional argument -- e.g. ``pplt.Cycle('Blues', 5)``. Using e.g.
# ``ax.plot(data, cycle='Blues')`` will automatically use the same
# number of samples as the number of columns in the dataset.
#
# Positional arguments passed to `~proplot.constructor.Cycle` are interpreted
# by the `~proplot.constructor.Colormap` constructor function, and the resulting
# colormap is sampled at discrete values. To exclude near-white colors on the
# end of a colormap, pass e.g. ``left=x`` to `~proplot.constructor.Cycle`, or
# supply a plotting command with e.g. ``cycle_kw={'left': x}``. See
# the :ref:`colormaps section <ug_cmaps>` for details.
#
# In the below example, several cycles are constructed from scratch, and the
# lines are referenced with colorbars and legends. Note that proplot allows
# you to :ref:`generate colorbars from lists of artists <ug_cbars>`.

# %%
import proplot as pplt
import numpy as np
fig = pplt.figure(refwidth=2, share=False)
state = np.random.RandomState(51423)
data = (20 * state.rand(10, 21) - 10).cumsum(axis=0)

# Cycle from on-the-fly monochromatic colormap
ax = fig.subplot(121)
lines = ax.plot(data[:, :5], cycle='plum', lw=5)
fig.colorbar(lines, loc='b', col=1, values=np.arange(0, len(lines)))
fig.legend(lines, loc='b', col=1, labels=np.arange(0, len(lines)))
ax.format(title='Cycle from a single color')

# Cycle from registered colormaps
ax = fig.subplot(122)
cycle = pplt.Cycle('blues', 'reds', 'oranges', 15, left=0.1)
lines = ax.plot(data[:, :15], cycle=cycle, lw=5)
fig.colorbar(lines, loc='b', col=2, values=np.arange(0, len(lines)), locator=2)
fig.legend(lines, loc='b', col=2, labels=np.arange(0, len(lines)), ncols=4)
ax.format(title='Cycle from merged colormaps', suptitle='Color cycles from colormaps')

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_cycles_other:
#
# Cycles of other properties
# --------------------------
#
# `~proplot.constructor.Cycle` can also generate cyclers that change
# properties other than color. Below, a single-color dash style cycler is
# constructed and applied to the axes locally. To apply it globally, simply
# use ``pplt.rc['axes.prop_cycle'] = cycle``.

# %%
import proplot as pplt
import numpy as np
import pandas as pd

# Cycle that loops through 'dashes' Line2D property
cycle = pplt.Cycle(lw=3, dashes=[(1, 0.5), (1, 1.5), (3, 0.5), (3, 1.5)])

# Sample data
state = np.random.RandomState(51423)
data = (state.rand(20, 4) - 0.5).cumsum(axis=0)
data = pd.DataFrame(data, columns=pd.Index(['a', 'b', 'c', 'd'], name='label'))

# Plot data
fig, ax = pplt.subplots(refwidth=2.5)
ax.format(suptitle='Plot without color cycle')
obj = ax.plot(
    data, cycle=cycle, legend='ll',
    legend_kw={'ncols': 2, 'handlelength': 2.5}
)

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_cycles_dl:
#
# Downloading color cycles
# ------------------------
#
# There are plenty of online interactive tools for generating and testing
# color cycles, including
# `i want hue <http://tools.medialab.sciences-po.fr/iwanthue/index.php>`__,
# `coolers <https://coolors.co>`__, and
# `viz palette <https://projects.susielu.com/viz-palette>`__.
#
# To add color cycles downloaded from any of these sources, add a cycle data file
# to the ``cycles`` subfolder inside `~proplot.config.Configurator.user_folder`
# and call `~proplot.config.register_cycles` (or restart your python session). You
# can also use `~proplot.colors.DiscreteColormap.from_file` or manually pass
# discrete colormaps or file paths to `~proplot.config.register_cycles`. See
# `~proplot.colors.DiscreteColormap.from_file` for a table of valid
# data file extensions.


text:
import cv2 as cv
import numpy as np

def to_rgb(image, color="BGR", max_val=None):
    if isinstance(image, str):
        image = cv.imread(image, 1)
        color = "BGR"

    assert isinstance(image, np.ndarray)

    if image.dtype != np.uint8:
        if max_val is None:
            min_val = image.min()
            max_val = image.max()
            scale = max_val - min_val
            if scale > 1e-3:
                image = (image - min_val) / scale
        else:
            image = image / max_val

        image = np.clip(image * 255, 0, 255)
        image = image.astype(np.uint8)

    if image.ndim == 2:
        image = cv.cvtColor(image, cv.COLOR_GRAY2RGB)
        color = "RGB"

    if color != "RGB":
        code = getattr(cv, f"COLOR_{color.upper()}2RGB")
        image = cv.cvtColor(image, code)

    return image

def crop_image(image, bbox):
    x1, y1, w, h = list(map(int, bbox))
    return image[y1:y1 + h, x1:x1 + w]

max_stars_repo_name: flystarhe/cvtk

max_issues_repo_name: flystarhe/cvtk

max_forks_repo_name: flystarhe/cvtk


text:
"""
The main file to run BSDE solver to solve parabolic partial differential equations (PDEs).

"""

import json
import logging
import os
import numpy as np

import torch
os.environ['KMP_DUPLICATE_LIB_OK']='True' # This is to fix libiomp5.dylib problem running OMP on mac

from config import get_config
from equation import get_equation
from solver import FeedForwardModel

import argparse
parser = argparse.ArgumentParser(description='Optional app description')
parser.add_argument('-pn', '--problem_name', type=str, default='HJB', help='The name of partial differential equation.')
parser.add_argument('-nr','--num_run', type=int, default=1, help='The number of experiments to repeatedly run for the same problem.')
parser.add_argument('-ld','--logdir', type=str, default='./logs', help='Directory where to write event logs and output array.')
args = parser.parse_args()

def main():
    problem_name = args.problem_name

    config = get_config(problem_name)
    bsde = get_equation(problem_name, config.dim, config.total_time, config.num_time_interval)

    if not os.path.exists(args.logdir):
        os.mkdir(args.logdir)

    path_prefix = os.path.join(args.logdir, problem_name)

    with open('{}_config.json'.format(path_prefix), 'w') as outfile:
        json.dump(dict((name, getattr(config, name))
                       for name in dir(config) if not name.startswith('__')),
                  outfile, indent=2)
    logging.basicConfig(level=logging.INFO,
                        format='%(levelname)-6s %(message)s')

    for idx_run in range(1, args.num_run+1):
        logging.info('Begin to solve %s with run %d' % (problem_name, idx_run))
        model = FeedForwardModel(config, bsde)
        if bsde.y_init:
            logging.info('Y0_true: %.4e' % bsde.y_init)
        training_history = model.train()
        if bsde.y_init:
            logging.info('relative error of Y0: %s',
                         '{:.2%}'.format(
                             abs(bsde.y_init - training_history[-1, 2])/bsde.y_init))
        # save training history
        np.savetxt('{}_training_history_{}.csv'.format(path_prefix, idx_run),
                   training_history,
                   fmt=['%d', '%.5e', '%.5e', '%d'],
                   delimiter=",",
                   header="step,loss_function,target_value,elapsed_time",
                   comments='')

if __name__ == '__main__':
    main()

max_stars_repo_name: XiaohanZhangCMU/spinningup

max_issues_repo_name: XiaohanZhangCMU/spinningup

max_forks_repo_name: XiaohanZhangCMU/spinningup


text:
import dace
import numpy as np
import re
from dace.memlet import Memlet
from dace import dtypes
from math import ceil

def add_cublas_cusolver(sdfg: dace.SDFG):
    """ Add CUBLAS and CUSOLVER handles to the SDFG. """
    sdfg.set_global_code("""
        #include <iostream>
        #include <cublas_v2.h>
        #include <cusolverDn.h>
        #include <cusparse.h>
        #include <thrust/complex.h>
        cublasHandle_t handle;
        float* const_zero;
        float* const_pone;
        """)
    sdfg.set_init_code("""
            cublasCreate(&handle);
            cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_DEVICE);
            cudaMalloc<float>(&const_zero, sizeof(float) *
            1);
            float zero = 0.0;
            cudaMemcpy(const_zero, &zero, sizeof(float) * 1,
            cudaMemcpyHostToDevice);
            cudaMalloc<float>(&const_pone, sizeof(float) *
            1);
            float pone = 1.0;
            cudaMemcpy(const_pone, &pone, sizeof(float) * 1,
            cudaMemcpyHostToDevice);
        """)
    sdfg.set_exit_code("""
            cublasDestroy(handle);
            cudaFree(const_zero);
            cudaFree(const_pone);
        """)

# TODO: Use library nodes

# Make sure C_memlet has a wcr if map_exit is used
# Matrix multiplication with one matrix as constant
def mm_small(
    state,
    A_node,
    B_node,
    C_node,
    A_subset=None,
    B_subset=None,
    C_subset=None,
    A_memlet=None,
    B_memlet=None,
    C_memlet=None,
    map_entry=None,
    map_exit=None,
    A_direct=True,
    B_direct=True,
):
    # C = A@B
    sdfg = state.parent
    Cshape = C_node.desc(sdfg).shape if C_subset is None else C_subset
    if isinstance(B_node, str):
        someshape = list(
            A_node.desc(sdfg).shape) if A_subset is None else A_subset
        someshape = [someshape[0]] + [someshape[1]] + [Cshape[1]]
    elif isinstance(A_node, str):
        someshape = list(
            B_node.desc(sdfg).shape) if B_subset is None else B_subset
        someshape = [Cshape[0]] + [someshape[0]] + [someshape[1]]
    else:
        raise AssertionError(
            "one of these have to be an str, else don't call this function")

    mapRange = ["0:" + str(_s) for _s in someshape]
    mapParams = ["i2", "i3", "i4"]
    mmapEntry, mmapExit = state.add_map(
        "matmul_sequential",
        dict(zip(mapParams, mapRange)),
        schedule=dace.ScheduleType.Sequential,
    )
    if isinstance(A_node, str):
        tasklet = state.add_tasklet("matmul_sequential", {"j1"}, {"out"},
                                    "out=" + A_node + "[i2, i3]" + "*j1")
        if B_memlet:
            state.add_edge(map_entry, None, mmapEntry, None, B_memlet)
            b_memlet_trailing = [str(_t[0]) for _t in B_memlet.subset[-2:]]
            state.add_edge(
                mmapEntry,
                None,
                tasklet,
                "j1",
                Memlet.simple(B_node,
                              ",".join(["i3", "i4"] + b_memlet_trailing)),
            )
        else:
            state.add_edge(
                B_node if B_direct else map_entry,
                None,
                mmapEntry,
                None,
                Memlet.from_array(B_node, B_node.desc(sdfg)),
            )
            state.add_edge(
                mmapEntry,
                None,
                tasklet,
                "j1",
                Memlet.simple(B_node, ",".join(["i3", "i4"])),
            )
    else:
        tasklet = state.add_tasklet("matmul_sequential", {"j0"}, {"out"},
                                    "out=j0*" + B_node + "[i3, i4]")
        if A_memlet:
            state.add_edge(map_entry, None, mmapEntry, None, A_memlet)
            a_memlet_trailing = [str(_t[0]) for _t in A_memlet.subset[-2:]]
            state.add_edge(
                mmapEntry,
                None,
                tasklet,
                "j0",
                Memlet.simple(A_node,
                              ",".join(["i2", "i3"] + a_memlet_trailing)),
            )
        else:
            state.add_edge(
                A_node if A_direct else map_entry,
                None,
                mmapEntry,
                None,
                Memlet.from_array(A_node, A_node.desc(sdfg)),
            )
            state.add_edge(
                mmapEntry,
                None,
                tasklet,
                "j0",
                Memlet.simple(A_node, ",".join(["i2", "i3"])),
            )

    if C_memlet:
        c_memlet_trailing = [str(_t[0]) for _t in C_memlet.subset[-2:]]
        state.add_edge(
            tasklet,
            "out",
            mmapExit,
            None,
            Memlet.simple(
                C_node,
                ",".join(["i2", "i4"] + c_memlet_trailing),
                wcr_str="lambda a,b: a+b",
                wcr_conflict=False,
            ),
        )
        state.add_edge(mmapExit, None, map_exit, None, C_memlet)
    else:
        state.add_edge(
            tasklet,
            "out",
            mmapExit,
            None,
            Memlet.simple(
                C_node,
                ",".join(["i2", "i4"]),
                wcr_str="lambda a,b:a+b",
                wcr_conflict=False,
            ),
        )
        state.add_edge(
            mmapExit,
            None,
            C_node if map_exit is None else map_exit,
            None,
            Memlet.simple(
                C_node,
                ",".join(["0:" + str(_s) for _s in C_node.desc(sdfg).shape]),
                wcr_str="lambda a,b:a+b",
                wcr_conflict=False,
            ),
        )

# takes input and stores output in column major order. give swapped input (B, A)
# instead of (A, B)
def mm(
    state,
    A_node,
    B_node,
    C_node,
    A_mode: str = "N",
    B_mode: str = "N",
    label: str = None,
    A_subset=None,
    B_subset=None,
    C_subset=None,
    A_memlet=None,
    B_memlet=None,
    C_memlet=None,
    map_entry=None,
    map_exit=None,
    shadow_a=False,
    shadow_b=False,
    buffer_a=False,
    buffer_c=False,
):
    sdfg = state.parent
    Adesc = A_node.desc(sdfg)
    Bdesc = B_node.desc(sdfg)
    Cdesc = C_node.desc(sdfg)
    Ashape = Adesc.shape if A_subset is None else A_subset
    Bshape = Bdesc.shape if B_subset is None else B_subset
    Cshape = Cdesc.shape if C_subset is None else C_subset

    # TODO change here for subset
    kdim_A = 0 if A_mode == "N" else 1
    lda = Ashape[1]
    ldb = Bshape[1]
    ldc = Cshape[1]

    # Set label
    if label is None:
        label = state.label

    # Create tasklet
    tasklet = state.add_tasklet(
        name=label + "_" + "mm_tasklet",
        inputs={"a", "b"},
        outputs={"c"},
        code="""
        cublasSetStream(handle, __dace_current_stream);
        cublasStatus_t status = cublasSgemm(
            handle,
            CUBLAS_OP_{amode}, CUBLAS_OP_{bmode},
            {m}, {n}, {k},
            const_pone,
            (float*)a, {lda},
            (float*)b, {ldb},
            const_zero,
            (float*)c, {ldc}
        );
        if (status)
            printf("Multiplication {a}*{b}->{c} failed (status %d)\\n", status);
        """.format(
            a=A_node.data,
            b=B_node.data,
            c=C_node.data,
            amode=A_mode,
            bmode=B_mode,
            m=Cshape[1],
            n=Cshape[0],
            k=Ashape[kdim_A],
            lda=lda,
            ldb=ldb,
            ldc=ldc,
        ),
        language=dace.dtypes.Language.CPP,
    )

    if not buffer_a:
        state.add_edge(
            A_node if not shadow_a else map_entry,
            None,
            tasklet,
            "a",
            A_memlet or dace.Memlet.from_array(A_node.data, Adesc),
        )
    else:
        # Assume that map_entry is to be used
        memcopy_tasklet = state.add_tasklet(label + "_" + "buffer_copy",
                                            {"j0"}, {"out"}, "out=j0")
        bufMapEntry, bufMapExit = state.add_map(
            label + "_bufmap",
            dict(zip(["i3", "i4"], ["0:" + str(_s) for _s in A_subset])),
            schedule=dace.ScheduleType.GPU_Device,
        )
        state.add_edge(map_entry, None, bufMapEntry, None, A_memlet)
        state.add_edge(
            bufMapEntry,
            None,
            memcopy_tasklet,
            "j0",
            Memlet.simple(
                A_node,
                ",".join(["i3", "i4"] +
                         [str(_t[0]) for _t in A_memlet.subset[-2:]]),
            ),
        )
        bufferNode = state.add_transient(
            A_node.data + "_buffercopy",
            A_subset,
            Adesc.dtype,
            storage=dace.StorageType.GPU_Global,
        )
        state.add_edge(memcopy_tasklet, "out", bufMapExit, None,
                       Memlet.simple(bufferNode, "i3,i4"))
        state.add_edge(
            bufMapExit,
            None,
            bufferNode,
            None,
            dace.Memlet.from_array(bufferNode, bufferNode.desc(sdfg)),
        )
        state.add_edge(
            bufferNode,
            None,
            tasklet,
            "a",
            dace.Memlet.from_array(bufferNode, bufferNode.desc(sdfg)),
        )
    state.add_edge(
        B_node if not shadow_b else map_entry,
        None,
        tasklet,
        "b",
        B_memlet or dace.Memlet.from_array(B_node.data, Bdesc),
    )
    if not buffer_c:
        state.add_edge(
            tasklet,
            "c",
            C_node if map_exit is None else map_exit,
            None,
            C_memlet or dace.Memlet.from_array(C_node.data, Cdesc),
        )
    else:
        memcopy_tasklet = state.add_tasklet(label + "_" + "buffer_copy_output",
                                            {"j0"}, {"out"}, "out=j0")
        bufMapEntry, bufMapExit = state.add_map(
            label + "_bufmap",
            dict(zip(["i3", "i4"], ["0:" + str(_s) for _s in C_subset])),
            schedule=dace.ScheduleType.GPU_Device,
        )
        bufferNode = state.add_transient(
            C_node.data + "_buffercopy_output",
            C_subset,
            Cdesc.dtype,
            storage=dace.StorageType.GPU_Global,
        )
        state.add_edge(
            tasklet,
            "c",
            bufferNode,
            None,
            dace.Memlet.from_array(bufferNode, bufferNode.desc(sdfg)),
        )
        state.add_edge(
            bufferNode,
            None,
            bufMapEntry,
            None,
            dace.Memlet.from_array(bufferNode, bufferNode.desc(sdfg)),
        )
        state.add_edge(bufMapEntry, None, memcopy_tasklet, "j0",
                       Memlet.simple(bufferNode, "i3,i4"))

        state.add_edge(
            memcopy_tasklet,
            "out",
            bufMapExit,
            None,
            Memlet.simple(
                C_node,
                ",".join(["i3", "i4"] +
                         [str(_t[0]) for _t in C_memlet.subset[-2:]]),
            ),
        )
        state.add_edge(bufMapExit, None, map_exit, None, C_memlet)

IMAGE_TILE_SIZE = 4
OUTPUT_TILE_SIZE = 2

bt = np.array([[1, 0, -1, 0], [0, 1, 1, 0], [0, -1, 1, 0],
               [0, 1, 0, -1]]).astype(np.float32)
b = np.array([[1, 0, 0, 0], [0, 1, -1, 1], [-1, 1, 1, 0],
              [0, 0, 0, -1]]).astype(np.float32)
g = np.array([[1.0, 0, 0], [0.5, 0.5, 0.5], [0.5, -0.5, 0.5],
              [0, 0, 1]]).astype(np.float32)
gt = np.array([[1, 0.5, 0.5, 0], [0, 0.5, -0.5, 0], [0, 0.5, 0.5,
                                                     1]]).astype(np.float32)
at = np.array([[1, 1, 1, 0], [0, 1, -1, -1]]).astype(np.float32)
a = np.array([[1, 0], [1, 1], [1, -1], [0, -1]]).astype(np.float32)

def printer(*inp):
    print(inp)

def string_builder(string):
    """ To match DaCe variable naming conventions, replaces all undesired
        characters with "_".
    """
    newstring = string
    if string[0].isdigit():
        newstring = "_" + string
    out = re.sub("[^a-zA-Z0-9_]", "_", newstring)
    return out

def winograd_convolution(dace_session, tf_node):
    debugNodes = []
    state = dace_session.state
    add_cublas_cusolver(dace_session.graph)
    #############Add constants for transformation matrices###############
    dace_session.graph.add_constant('Btrans', bt)
    dace_session.graph.add_constant('B', b)
    bNode = 'B'
    bTransposeNode = 'Btrans'

    dace_session.graph.add_constant('G', g)
    dace_session.graph.add_constant('Gtrans', gt)
    gNode = 'G'
    gTransposeNode = 'Gtrans'

    dace_session.graph.add_constant('Atrans', at)
    dace_session.graph.add_constant('A', a)
    aNode = 'A'
    aTransposeNode = 'Atrans'

    inputNodes = []
    inputParams = []
    inputDims = []
    for _inp in tf_node.inputs:
        _node, _params, _dims = dace_session.create_and_add_input_node(_inp)
        inputNodes.append(_node)
        inputParams.append(_params)
        inputDims.append(_dims)
    # Manually add copy for kernel from CPU to GPU
    kernel_desc = inputNodes[1].desc(dace_session.graph)
    kernelGPU = state.add_transient(
        inputNodes[1].data + "GPU",
        shape=kernel_desc.shape,
        dtype=kernel_desc.dtype,
        lifetime=dtypes.AllocationLifetime.SDFG,
        storage=dace.StorageType.GPU_Global,
    )
    state.add_edge(
        inputNodes[1],
        None,
        kernelGPU,
        None,
        Memlet.from_array(inputNodes[1],
                          inputNodes[1].desc(dace_session.graph)),
    )
    inputNodes[1] = kernelGPU
    outputList = dace_session.create_and_add_output_node(tf_node)
    outputDims = dace_session.get_default_dims(tf_node.outputs[0])
    if str(tf_node.get_attr("padding"))[2:-1] == "SAME":
        paddedInput, paddedDims = dace_session.inputPadding(
            tf_node,
            inputNodes[0],
            inputNodes[0].desc(dace_session.graph),
            outputList[0].desc(dace_session.graph).shape[1],
            inputNodes[1].desc(dace_session.graph).shape[0],
            tf_node.get_attr("strides")[1],
            inputDims[0],
        )
        inputDims[0] = paddedDims
        inputNodes[0] = paddedInput
    outputShape = [int(_s) for _s in tf_node.outputs[0].shape]
    inputViewShape = [
        IMAGE_TILE_SIZE,
        IMAGE_TILE_SIZE,
        tf_node.inputs[0].shape[-1],
        outputShape[0] * ceil(outputShape[1] / OUTPUT_TILE_SIZE) *
        ceil(outputShape[2] / OUTPUT_TILE_SIZE),
    ]
    inputViewDims = ["0:" + str(_x) for _x in inputViewShape]
    ########Tiling the image#################################
    inputViewParams = [
        "i3%" + str(outputShape[0]),
        "(i3/" + str(outputShape[0]) + ")%"
        # + str(output_shape[0] * ceil(output_shape[1] / OUTPUT_TILE_SIZE))
        + str(ceil(outputShape[2] / OUTPUT_TILE_SIZE)) + "*" +
        str(OUTPUT_TILE_SIZE) + "+i0",
        # + str(
        #    ceil(output_shape[1] / OUTPUT_TILE_SIZE)
        #    * ceil(output_shape[2] / OUTPUT_TILE_SIZE)
        # ),
        "int_floor(i3,"
        # + str(ceil(output_shape[1] / OUTPUT_TILE_SIZE))
        + str(outputShape[0] * ceil(outputShape[2] / OUTPUT_TILE_SIZE)) +
        ")*" + str(OUTPUT_TILE_SIZE) + "+i1",
        "i2",
    ]
    inputView = state.add_transient(
        "V" + "_".join([str(_s) for _s in inputViewShape]),
        inputViewShape,
        dace.float32,
        dace.StorageType.GPU_Global,
    )
    mapEntry, mapExit = state.add_map(
        string_builder(tf_node.name) + "_input_tile",
        dict(zip(inputParams[0], inputViewDims)),
    )
    tasklet = state.add_tasklet(
        string_builder(tf_node.name) + "_input_tile", {"j0"}, {"out"},
        "out = j0")
    dace_session.add_in_memlets([inputNodes[0]], mapEntry, tasklet,
                                [inputDims[0]], [inputViewParams])
    dace_session.add_out_memlets([inputView], mapExit, tasklet,
                                 [inputViewDims], [inputParams[0]])
    ##################Transforming all input tiles#########################
    #[TODO] try to re-use memory
    vNode = state.add_transient(
        "V_output" + "_".join([str(_s) for _s in inputViewShape]),
        inputViewShape,
        dace.float32,
        dace.StorageType.GPU_Global,
    )
    vNode.setzero = True
    mapEntry, mapExit = state.add_map(
        string_builder(tf_node.name) + "_input_txform",
        dict(zip(inputParams[0][0:2], inputViewDims[2:4])),
        dace.ScheduleType.GPU_Device,
    )
    intermediateResultNode = state.add_transient("BtI", bt.shape, dace.float32,
                                                 dace.StorageType.Register)
    intermediateResultNode.setzero = True
    state.add_edge(
        inputView,
        None,
        mapEntry,
        None,
        Memlet.simple(inputView, ",".join(inputViewDims)),
    )
    mm_small(
        state,
        bTransposeNode,
        inputView,
        intermediateResultNode,
        B_subset=[IMAGE_TILE_SIZE, IMAGE_TILE_SIZE],
        B_memlet=Memlet.simple(
            inputView, ",".join(inputViewDims[0:2] + inputParams[0][0:2])),
        map_entry=mapEntry,
        B_direct=False,
    )
    mm_small(
        state,
        intermediateResultNode,
        bNode,
        vNode,
        map_exit=mapExit,
        C_subset=[IMAGE_TILE_SIZE, IMAGE_TILE_SIZE],
        C_memlet=Memlet.simple(
            vNode,
            ",".join(inputViewDims[0:2] + inputParams[0][0:2]),
            wcr_str="lambda a,b: a+b",
            wcr_conflict=False,
        ),
        map_entry=mapEntry,
        A_direct=True,
    )
    state.add_edge(
        mapExit,
        None,
        vNode,
        None,
        Memlet.simple(
            vNode,
            ",".join(inputViewDims),
            wcr_str="lambda a,b: a+b",
            wcr_conflict=False,
        ),
    )
    #############Transforming the kernel###############################
    mapEntry, mapExit = state.add_map(
        string_builder(tf_node.name) + "_kernel_txform",
        dict(zip(inputParams[1][0:2], inputDims[1][2:4])),
        dace.ScheduleType.GPU_Device,
    )
    intermediateResultNode = state.add_transient("GF", g.shape, dace.float32,
                                                 dace.StorageType.Register)
    intermediateResultNode.setzero = True
    processedKernelNode = state.add_transient(
        "U" + "_".join([
            str(_s) for _s in inputViewShape[0:2] +
            list(tf_node.inputs[1].shape[-1:-3:-1])
        ]),
        inputViewShape[0:2] + list(tf_node.inputs[1].shape[-1:-3:-1]),
        dace.float32,
        dace.StorageType.GPU_Global,
    )
    processedKernelNode.setzero = True
    state.add_edge(
        inputNodes[1],
        None,
        mapEntry,
        None,
        dace.Memlet.from_array(inputNodes[1].data,
                               inputNodes[1].desc(dace_session.graph)),
    )

    mm_small(
        state,
        gNode,
        inputNodes[1],
        intermediateResultNode,
        map_entry=mapEntry,
        B_subset=tf_node.inputs[1].shape[0:2],
        B_memlet=Memlet.simple(
            inputNodes[1], ",".join(inputDims[1][0:2] + inputParams[1][0:2])),
        B_direct=False,
    )
    mm_small(
        state,
        intermediateResultNode,
        gTransposeNode,
        processedKernelNode,
        C_subset=[IMAGE_TILE_SIZE, IMAGE_TILE_SIZE],
        C_memlet=Memlet.simple(
            processedKernelNode,
            ",".join(inputViewDims[0:2] + [inputParams[0][1]] +
                     [inputParams[0][0]]),
            wcr_str="lambda a,b: a+b",
            wcr_conflict=False,
        ),
        map_entry=mapEntry,
        map_exit=mapExit,
        A_direct=True,
    )
    state.add_edge(
        mapExit,
        None,
        processedKernelNode,
        None,
        Memlet.simple(
            processedKernelNode.data,
            ",".join([
                "0:" + str(_s)
                for _s in processedKernelNode.desc(dace_session.graph).shape
            ]),
            wcr_str="lambda a,b: a+b",
            wcr_conflict=False,
        ),
    )
    mNode = state.add_transient(
        "m" + "_".join([
            str(_s) for _s in inputViewShape[0:2] +
            [tf_node.inputs[1].shape[-1], inputViewShape[-1]]
        ]),
        inputViewShape[0:2] +
        [tf_node.inputs[1].shape[-1], inputViewShape[-1]],
        dace.float32,
        dace.StorageType.GPU_Global,
    )
    mNodeDims = ["0:" + str(_d) for _d in mNode.desc(dace_session.graph).shape]
    mapEntry, mapExit = state.add_map(
        string_builder(tf_node.name) + "_eltwise_product",
        dict(zip(inputParams[0][0:2], inputViewDims[0:2])),
        dace.ScheduleType.Sequential,
    )
    state.add_edge(
        vNode,
        None,
        mapEntry,
        None,
        Memlet.from_array(vNode.data, vNode.desc(dace_session.graph)),
    )
    state.add_edge(
        processedKernelNode,
        None,
        mapEntry,
        None,
        Memlet.from_array(processedKernelNode.data,
                          processedKernelNode.desc(dace_session.graph)),
    )
    mm(
        state,
        vNode,
        processedKernelNode,
        mNode,
        A_subset=inputViewShape[2:4],
        A_memlet=Memlet.simple(
            vNode, ",".join(inputParams[0][0:2] + inputViewDims[-2:])),
        B_subset=tf_node.inputs[1].shape[-1:-3:-1],
        B_memlet=Memlet.simple(
            processedKernelNode,
            ",".join(
                inputParams[0][0:2] +
                ["0:" + str(_s) for _s in tf_node.inputs[1].shape[-1:-3:-1]]),
        ),
        C_subset=[tf_node.inputs[1].shape[-1], inputViewShape[-1]],
        C_memlet=Memlet.simple(mNode,
                               ",".join(inputParams[0][0:2] + mNodeDims[-2:])),
        map_entry=mapEntry,
        map_exit=mapExit,
        shadow_a=True,
        shadow_b=True,
    )
    state.add_edge(mapExit, None, mNode, None,
                   Memlet.simple(mNode, ",".join(mNodeDims)))
    #################OUTPUT TRANSFORMATION################################
    mapRange = [inputDims[1][-1]] + [inputViewDims[-1]]
    mapEntry, mapExit = state.add_map(
        string_builder(tf_node.name) + "_output_txform",
        dict(zip(inputParams[0][0:2], mapRange)),
        dace.ScheduleType.GPU_Device,
    )
    intermediateResultNode = state.add_transient("AtM", at.shape, dace.float32,
                                                 dace.StorageType.Register)
    intermediateResultNode.setzero = True
    transformedOutputNode = state.add_transient(
        "inv_txformed_output" + "_".join([str(tf_node.inputs[1].shape[-1])] +
                                         [str(inputViewShape[-1])]),
        [OUTPUT_TILE_SIZE, OUTPUT_TILE_SIZE] + [tf_node.inputs[1].shape[-1]] +
        [inputViewShape[-1]],
        dace.float32,
        dace.StorageType.GPU_Global,
    )
    transformedOutputNode.setzero = True
    state.add_edge(mNode, None, mapEntry, None,
                   Memlet.simple(mNode, ",".join(mNodeDims)))
    mm_small(
        state,
        aTransposeNode,
        mNode,
        intermediateResultNode,
        B_subset=inputViewShape[0:2],
        B_memlet=Memlet.simple(
            mNode, ",".join(inputViewDims[0:2] + inputParams[0][0:2])),
        map_entry=mapEntry,
        B_direct=False,
    )
    mm_small(
        state,
        intermediateResultNode,
        aNode,
        transformedOutputNode,
        C_subset=[OUTPUT_TILE_SIZE, OUTPUT_TILE_SIZE],
        C_memlet=Memlet.simple(
            transformedOutputNode,
            ",".join(
                ["0:" + str(OUTPUT_TILE_SIZE), "0:" + str(OUTPUT_TILE_SIZE)] +
                inputParams[0][0:2]),
            wcr_str="lambda a,b:a+b",
            wcr_conflict=False,
        ),
        map_entry=mapEntry,
        map_exit=mapExit,
        A_direct=True,
    )
    state.add_edge(
        mapExit,
        None,
        transformedOutputNode,
        None,
        Memlet.simple(
            transformedOutputNode.data,
            ",".join([
                "0:" + str(_s)
                for _s in transformedOutputNode.desc(dace_session.graph).shape
            ]),
            wcr_str="lambda a,b: a+b",
            wcr_conflict=False,
        ),
    )
    ###################Un-Tile the output to NHWC format###################
    outputParams = [
        "i3%" + str(outputShape[0]),
        "(i3/" + str(outputShape[0]) + ")%" +
        str(ceil(outputShape[2] / OUTPUT_TILE_SIZE)) + "*" +
        str(OUTPUT_TILE_SIZE) + "+i0",
        "int_floor(i3," +
        str(outputShape[0] * ceil(outputShape[2] / OUTPUT_TILE_SIZE)) + ")*" +
        str(OUTPUT_TILE_SIZE) + "+i1",
        "i2",
    ]
    mapRange = [
        "0:" + str(_s)
        for _s in transformedOutputNode.desc(dace_session.graph).shape
    ]
    mapEntry, mapExit = state.add_map(
        string_builder(tf_node.name) + "_output_untile",
        dict(zip(inputParams[0], mapRange)),
    )
    tasklet = state.add_tasklet(
        string_builder(tf_node.name) + "_output_untile", {"j0"}, {"out"},
        "out = j0")
    dace_session.add_in_memlets([transformedOutputNode], mapEntry, tasklet,
                                [mapRange], [inputParams[0]])
    dace_session.add_out_memlets(outputList, mapExit, tasklet, [outputDims],
                                 [outputParams])
    ################# Debugging with callbacks #############
    taskletInputs = ["i" + str(index) for index in range(len(debugNodes))]
    callback_tasklet = state.add_tasklet(
        string_builder(tf_node.name) + "_printer",
        {*taskletInputs},
        {},
        string_builder(tf_node.name) + "_printer" + "(" +
        ",".join(taskletInputs) + ");",
        language=dace.dtypes.Language.CPP,
    )
    for _n, _conn in zip(debugNodes, taskletInputs):
        _n_cpu = state.add_transient(_n.data + "_cpucopy",
                                     _n.desc(dace_session.graph).shape,
                                     _n.desc(dace_session.graph).dtype,
                                     storage=dace.StorageType.CPU_Heap,
                                     lifetime=dtypes.AllocationLifetime.SDFG)
        state.add_edge(_n, None, _n_cpu, None,
                       Memlet.from_array(_n, _n.desc(dace_session.graph)))
        state.add_edge(
            _n_cpu,
            None,
            callback_tasklet,
            _conn,
            Memlet.from_array(_n_cpu, _n_cpu.desc(dace_session.graph)),
        )
    callback_input_types = []
    for somenode in debugNodes:
        callback_input_types.append(somenode.desc(dace_session.graph))
    dace_session.callbackFunctionDict[string_builder(tf_node.name) +
                                      "_printer"] = printer
    dace_session.callbackTypeDict[string_builder(tf_node.name) +
                                  "_printer"] = dace.data.Scalar(
                                      dace.callback(None,
                                                    *callback_input_types))

max_stars_repo_name: am-ivanov/dace


text:
import logging
from typing import cast, List, Union

import numpy

from cephlib.common import float2str
from cephlib.texttable import Texttable
from cephlib.statistic import calc_norm_stat_props, calc_histo_stat_props

from .stage import Stage, StepOrder
from .test_run_class import TestRun
from .result_classes import SuiteConfig
from .suits.io.fio import FioTest
from .suits.io.fio_job import FioJobParams
from .suits.io.fio_hist import get_lat_vals
from .data_selectors import get_aggregated
from .result_storage import IWallyStorage

logger = logging.getLogger("wally")

console_report_headers = ["Description", "IOPS ~ Dev", "BW, MiBps", 'Skew/Kurt', 'lat med, ms', 'lat 95, ms']
console_report_align = ['l', 'r', 'r', 'r', 'r', 'r']

def get_console_report_table(suite: SuiteConfig, rstorage: IWallyStorage) -> List[Union[List[str], Texttable.HLINE]]:
    table: List[Union[List[str], Texttable.HLINE]] = []
    prev_params = None
    for job in sorted(rstorage.iter_job(suite), key=lambda job: job.params):
        fparams = cast(FioJobParams, job.params)
        fparams['qd'] = None

        if prev_params is not None and fparams.char_tpl != prev_params:
            table.append(Texttable.HLINE)

        prev_params = fparams.char_tpl

        bw_ts = get_aggregated(rstorage, suite.storage_id, job.storage_id, metric='bw',
                               trange=job.reliable_info_range_s)
        props = calc_norm_stat_props(bw_ts)
        avg_iops = props.average // job.params.params['bsize']
        iops_dev = props.deviation // job.params.params['bsize']

        lat_ts = get_aggregated(rstorage, suite.storage_id, job.storage_id, metric='lat',
                                trange=job.reliable_info_range_s)
        bins_edges = numpy.array(get_lat_vals(lat_ts.data.shape[1]), dtype='float32') / 1000  # convert us to ms
        lat_props = calc_histo_stat_props(lat_ts, bins_edges)
        table.append([job.params.summary,
                      f"{float2str(avg_iops):>6s} ~ {float2str(iops_dev):>6s}",
                      float2str(props.average / 1024),  # Ki -> Mi
                      f"{props.skew:>5.1f}/{props.kurt:>5.1f}",
                      float2str(lat_props.perc_50), float2str(lat_props.perc_95)])
    return table

class ConsoleReportStage(Stage):

    priority = StepOrder.REPORT

    def run(self, ctx: TestRun) -> None:
        for suite in ctx.rstorage.iter_suite(FioTest.name):
            table = Texttable(max_width=200)
            table.set_deco(Texttable.VLINES | Texttable.BORDER | Texttable.HEADER)
            tbl = ctx.rstorage.get_txt_report(suite)
            if tbl is None:
                table.header(console_report_headers)
                table.set_cols_align(console_report_align)
                for line in get_console_report_table(suite, ctx.rstorage):
                    table.add_row(line)
                tbl = table.draw()
                ctx.rstorage.put_txt_report(suite, tbl)
            print(tbl)

max_stars_repo_name: Mirantis/rally-results-processor

max_issues_repo_name: Mirantis/rally-results-processor


text:
import numpy as np
import scipy.stats as st
from sklearn.metrics import accuracy_score, recall_score, precision_score, f1_score
from typing import Optional, List, Dict

class gmmhmm:
    # This class converted with modifications from https://code.google.com/p/hmm-speech-recognition/source/browse/Word.m
    def __init__(self, n_states, n_iter=15):
        self.n_states = n_states
        self.n_iter = n_iter
        self.random_state = np.random.RandomState(0)

        # Normalize random initial state
        self.prior = self._normalize(self.random_state.rand(self.n_states, 1))
        self.A = self._stochasticize(self.random_state.rand(self.n_states, self.n_states))

        self.mu = None
        self.covs = None
        self.n_dims = None

    def _forward(self, B):
        log_likelihood = 0.
        T = B.shape[1]
        alpha = np.zeros(B.shape)
        for t in range(T):
            if t == 0:
                alpha[:, t] = B[:, t] * self.prior.ravel()
            else:
                alpha[:, t] = B[:, t] * np.dot(self.A.T, alpha[:, t - 1])
            alpha[:, t] = alpha[:, t] + (alpha[:, t] == 0)
            alpha_sum = np.sum(alpha[:, t])
            alpha_sum = np.where(alpha_sum, alpha_sum, t)
            alpha[:, t] /= alpha_sum
            log_likelihood = log_likelihood + np.log(alpha_sum)
        return log_likelihood, alpha

    def _backward(self, B):
        T = B.shape[1]
        beta = np.zeros(B.shape)

        beta[:, -1] = np.ones(B.shape[0])

        for t in range(T - 1)[::-1]:
            beta[:, t] = np.dot(self.A, (B[:, t + 1] * beta[:, t + 1]))
            beta[:, t] = beta[:, t] + (beta[:, t] == 0)
            beta_sum = np.sum(beta[:, t])
            beta[:, t] /= np.where(beta_sum, beta_sum, t)
        return beta

    def _state_likelihood(self, obs):
        obs = np.atleast_2d(obs)
        B = np.zeros((self.n_states, obs.shape[1]))
        for s in range(self.n_states):
            # Needs scipy 0.14
            np.random.seed(self.random_state.randint(1))
            B[s, :] = st.multivariate_normal.pdf(
                obs.T, mean=self.mu[:, s].T, cov=self.covs[:, :, s].T)
            # This function can (and will!) return values >> 1
            # See the discussion here for the equivalent matlab function
            # https://groups.google.com/forum/#!topic/comp.soft-sys.matlab/YksWK0T74Ak
            # Key line: "Probabilities have to be less than 1,
            # Densities can be anything, even infinite (at individual points)."
            # This is evaluating the density at individual points...
        return B

    def _normalize(self, x):
        return (x + (x == 0)) / np.where(np.sum(x), np.sum(x), len(x))

    def _stochasticize(self, x):
        # robi to samo co _normalize tylko dla macierzy przejścia
        sum_of_rows = x.sum(axis=1)
        sum_of_rows = np.where(sum_of_rows, sum_of_rows, len(x))
        return (x + (x == 0)) / sum_of_rows[:, np.newaxis]

    def _em_init(self, obs):
        # Using this _em_init function allows for less required constructor args
        if self.n_dims is None:
            self.n_dims = obs.shape[0]
        if self.mu is None:
            subset = self.random_state.choice(np.arange(self.n_dims), size=self.n_states, replace=False)
            self.mu = obs[:, subset]
        if self.covs is None:
            self.covs = np.zeros((self.n_dims, self.n_dims, self.n_states))
            self.covs += np.diag(np.diag(np.cov(obs)))[:, :, None]
        return self

    def _em_step(self, obs):
        obs = np.atleast_2d(obs)
        B = self._state_likelihood(obs)
        T = obs.shape[1]

        log_likelihood, alpha = self._forward(B)
        beta = self._backward(B)

        xi_sum = np.zeros((self.n_states, self.n_states))
        gamma = np.zeros((self.n_states, T))

        for t in range(T - 1):
            partial_sum = self.A * np.dot(alpha[:, t], (beta[:, t] * B[:, t + 1]).T)
            xi_sum += self._normalize(partial_sum)
            partial_g = alpha[:, t] * beta[:, t]
            gamma[:, t] = self._normalize(partial_g)

        partial_g = alpha[:, -1] * beta[:, -1]
        gamma[:, -1] = self._normalize(partial_g)

        expected_prior = gamma[:, 0]
        expected_A = self._stochasticize(xi_sum)

        expected_mu = np.zeros((self.n_dims, self.n_states))
        expected_covs = np.zeros((self.n_dims, self.n_dims, self.n_states))

        gamma_state_sum = np.sum(gamma, axis=1)
        # Set zeros to 1 before dividing
        gamma_state_sum = gamma_state_sum + (gamma_state_sum == 0)

        for s in range(self.n_states):
            gamma_obs = obs * gamma[s, :]
            expected_mu[:, s] = np.sum(gamma_obs, axis=1) / gamma_state_sum[s]
            partial_covs = np.dot(gamma_obs, obs.T) / gamma_state_sum[s] - np.dot(expected_mu[:, s],
                                                                                expected_mu[:, s].T)
            # Symmetrize
            partial_covs = np.triu(partial_covs) + np.triu(partial_covs).T - np.diag(partial_covs)

        # Ensure positive semidefinite by adding diagonal loading
        expected_covs += .01 * np.eye(self.n_dims)[:, :, None]

        self.prior = expected_prior
        self.mu = expected_mu
        self.covs = expected_covs
        self.A = expected_A
        return log_likelihood

    def fit(self, obs):
        # Support for 2D and 3D arrays
        # 2D should be n_features, n_dims
        # 3D should be n_examples, n_features, n_dims
        # For example, with 6 features per speech segment, 105 different words
        # this array should be size
        # (105, 6, X) where X is the number of frames with features extracted
        # For a single example file, the array should be size (6, X)
        if len(obs.shape) == 2:
            for i in range(self.n_iter):
                self._em_init(obs)
                log_likelihood = self._em_step(obs)
        elif len(obs.shape) == 3:
            count = obs.shape[0]
            for n in range(count):
                for i in range(self.n_iter):
                    self._em_init(obs[n, :, :])
                    log_likelihood = self._em_step(obs[n, :, :])
        return self

    def transform(self, obs):
        # Support for 2D and 3D arrays
        # 2D should be n_features, n_dims
        # 3D should be n_examples, n_features, n_dims
        # For example, with 6 features per speech segment, 105 different words
        # this array should be size
        # (105, 6, X) where X is the number of frames with features extracted
        # For a single example file, the array should be size (6, X)
        if len(obs.shape) == 2:
            B = self._state_likelihood(obs)
            log_likelihood, _ = self._forward(B)
            return log_likelihood
        elif len(obs.shape) == 3:
            count = obs.shape[0]
            out = np.zeros((count,))
            for n in range(count):
                B = self._state_likelihood(obs[n, :, :])
                log_likelihood, _ = self._forward(B)
                out[n] = log_likelihood
            return out

class IwrGaussianHMMModel:
    """
        A class used to create Isolated Word Recognition dataset with Google Text to Speech tool

        ...

        Attributes
        ----------
        n_states : int
            number of hidden states in hmm model
        n_iter : int
            number of iterations in Baum-Welch algorithm
        model : List[gmmhmm, ...]
            list of hmm models, each model represents chain of single word class

        Methods
        -------
        fit
            fitting list of models
        _divide_dataset_by_class
            creating list of datasets for each word class
        predict
            predicting word class label
        score
            calculating solution score
        """

    def __init__(self, n_states: int, n_iter: Optional[int] = 10) -> None:
        """

        Parameters
        ----------
        n_states (int): number of hidden states in hmm
        n_iter (int): number of iterations in Baum-Welch algorithm
        """
        self.n_states = n_states
        self.n_iter = n_iter
        self.model = None

    def fit(self, X: np.array, y: np.array) -> None:
        """
        fitting list of models

        Parameters
        ----------
        X (np.array): array of shape (n_observations, n_mfcc_coef, n_windows)
        y (np.array): array of labels for each word class

        Returns
        -------
            None
        """
        # Unique labels
        labels = list(set(y))
        # list of datasets for each unique class
        datasets = self._divide_dataset_by_class(X, y, labels)

        # list of hmm models, each model represents one word class
        self.model = [gmmhmm(n_states=self.n_states,
                             n_iter=self.n_iter)
                      for _ in range(len(labels))]
        # fitting each model
        self.model = [m.fit(x) for m, x in zip(self.model, datasets)]

    def _divide_dataset_by_class(self, X: np.array, y: np.array, labels: List[float]) -> List[np.array]:
        """
        creating list of datasets for each word class

        Parameters
        ----------
        X (np.array): array of shape (n_observations, n_mfcc_coef, n_windows)
        y (np.array): array of labels for each word class
        labels (list): list of unique labels

        Returns
        -------
            datasets (list): list of array, each array contain single word class signal
        """
        datasets = [X[y == label, :, :] for label in labels]
        return datasets

    def predict(self, X: np.array) -> np.array:
        """
        predicting word class label

        Parameters
        ----------
        X (np.array): array of shape (n_observations, n_mfcc_coef, n_windows)

        Returns
        -------
        preds (np.array): array with prediction label
        """
        n, *_ = X.shape
        # if model exists
        if self.model:
            # each log likelihood values from models
            ps = [m.transform(X) for m in self.model]
            # stacking observations
            res = np.vstack(ps)
            # taking the max index
            preds = np.argmax(res, axis=0)
            return preds
        else:
            print('First of all you must fit a model')

    def score(self, X: np.array, y: np.array) -> Dict[str, float]:
        """
        calculating solution score

        Parameters
        ----------
        X (np.array): array of shape (n_observations, n_mfcc_coef, n_windows)
        y (np.array): array of labels for each word class

        Returns
        -------
            acc (float): accuracy of solutions
        """
        results = {}
        preds = self.predict(X)
        accuracy = accuracy_score(y, preds)
        recall = recall_score(y, preds, average='weighted', zero_division=1)
        precision = precision_score(y, preds, average='weighted', zero_division=1)
        f1 = f1_score(y, preds, average='weighted', zero_division=1)

        results['accuracy'] = accuracy
        results['recall'] = recall
        results['precision'] = precision
        results['f1'] = f1
        return results

reason: import numpy,import scipy


text:
#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''Density expansion on plane waves'''

import time
import ctypes
import numpy
import copy
from pyscf import lib
from pyscf import gto
from pyscf.lib import logger
from pyscf.pbc import tools
from pyscf.pbc.gto import pseudo
from pyscf.pbc.df import ft_ao
from pyscf.pbc.df import incore
from pyscf.pbc.df.df_jk import KPT_DIFF_TOL, is_zero, gamma_point
from pyscf.pbc.df import aft_jk
from pyscf.pbc.df import aft_ao2mo

def estimate_eta(cell, cutoff=1e-12):
    '''The exponent of the smooth gaussian model density, requiring that at
    boundary, density ~ 4pi rmax^2 exp(-eta/2*rmax^2) ~ 1e-12
    '''
    # r^5 to guarantee at least up to f shell converging at boundary
    eta = max(numpy.log(4*numpy.pi*cell.rcut**5/cutoff)/cell.rcut**2*2, .3)
    return eta

def get_nuc(mydf, kpts=None):
    cell = mydf.cell
    if kpts is None:
        kpts_lst = numpy.zeros((1,3))
    else:
        kpts_lst = numpy.reshape(kpts, (-1,3))

    log = logger.Logger(mydf.stdout, mydf.verbose)
    t1 = t0 = (time.clock(), time.time())

    nkpts = len(kpts_lst)
    nao = cell.nao_nr()
    nao_pair = nao * (nao+1) // 2

    Gv, Gvbase, kws = cell.get_Gv_weights(mydf.gs)
    kpt_allow = numpy.zeros(3)
    if mydf.eta == 0:
        vpplocG = pseudo.pp_int.get_gth_vlocG_part1(cell, Gv)
        vpplocG = -numpy.einsum('ij,ij->j', cell.get_SI(Gv), vpplocG)
        vpplocG *= kws
        vGR = vpplocG.real
        vGI = vpplocG.imag
        vjR = numpy.zeros((nkpts,nao_pair))
        vjI = numpy.zeros((nkpts,nao_pair))
    else:
        nuccell = copy.copy(cell)
        half_sph_norm = .5/numpy.sqrt(numpy.pi)
        norm = half_sph_norm/gto.mole._gaussian_int(2, mydf.eta)
        chg_env = [mydf.eta, norm]
        ptr_eta = cell._env.size
        ptr_norm = ptr_eta + 1
        chg_bas = [[ia, 0, 1, 1, 0, ptr_eta, ptr_norm, 0] for ia in range(cell.natm)]
        nuccell._atm = cell._atm
        nuccell._bas = numpy.asarray(chg_bas, dtype=numpy.int32)
        nuccell._env = numpy.hstack((cell._env, chg_env))

        # PP-loc part1 is handled by fakenuc in _int_nuc_vloc
        vj = lib.asarray(mydf._int_nuc_vloc(nuccell, kpts_lst))
        vjR = vj.real
        vjI = vj.imag
        t1 = log.timer_debug1('vnuc pass1: analytic int', *t1)

        charge = -cell.atom_charges()
        coulG = tools.get_coulG(cell, kpt_allow, gs=mydf.gs, Gv=Gv)
        coulG *= kws
        aoaux = ft_ao.ft_ao(nuccell, Gv)
        vGR = numpy.einsum('i,xi->x', charge, aoaux.real) * coulG
        vGI = numpy.einsum('i,xi->x', charge, aoaux.imag) * coulG

    max_memory = max(2000, mydf.max_memory-lib.current_memory()[0])
    for k, pqkR, pqkI, p0, p1 \
            in mydf.ft_loop(mydf.gs, kpt_allow, kpts_lst,
                            max_memory=max_memory, aosym='s2'):
# rho_ij(G) nuc(-G) / G^2
# = [Re(rho_ij(G)) + Im(rho_ij(G))*1j] [Re(nuc(G)) - Im(nuc(G))*1j] / G^2
        if not gamma_point(kpts_lst[k]):
            vjI[k] += numpy.einsum('k,xk->x', vGR[p0:p1], pqkI)
            vjI[k] -= numpy.einsum('k,xk->x', vGI[p0:p1], pqkR)
        vjR[k] += numpy.einsum('k,xk->x', vGR[p0:p1], pqkR)
        vjR[k] += numpy.einsum('k,xk->x', vGI[p0:p1], pqkI)
    t1 = log.timer_debug1('contracting Vnuc', *t1)

    vj = []
    for k, kpt in enumerate(kpts_lst):
        if gamma_point(kpt):
            vj.append(lib.unpack_tril(vjR[k]))
        else:
            vj.append(lib.unpack_tril(vjR[k]+vjI[k]*1j))

    if kpts is None or numpy.shape(kpts) == (3,):
        vj = vj[0]
    return vj

def _int_nuc_vloc(mydf, nuccell, kpts, intor='cint3c2e_sph'):
    '''Vnuc - Vloc'''
    cell = mydf.cell
    rcut = max(cell.rcut, nuccell.rcut)
    Ls = cell.get_lattice_Ls(rcut=rcut)
    expLk = numpy.asarray(numpy.exp(1j*numpy.dot(Ls, kpts.T)), order='C')
    nkpts = len(kpts)

# Use the 3c2e code with steep s gaussians to mimic nuclear density
    fakenuc = _fake_nuc(cell)
    fakenuc._atm, fakenuc._bas, fakenuc._env = \
            gto.conc_env(nuccell._atm, nuccell._bas, nuccell._env,
                         fakenuc._atm, fakenuc._bas, fakenuc._env)

    nao = cell.nao_nr()
    buf = [numpy.zeros((nao,nao,fakenuc.natm), order='F', dtype=numpy.complex128)
           for k in range(nkpts)]
    ints = incore._wrap_int3c(cell, fakenuc, intor, 1, Ls, buf)
    atm, bas, env = ints._envs[:3]
    c_shls_slice = (ctypes.c_int*6)(0, cell.nbas, cell.nbas, cell.nbas*2,
                                    cell.nbas*2, cell.nbas*2+fakenuc.natm)

    xyz = numpy.asarray(cell.atom_coords(), order='C')
    ptr_coordL = atm[:cell.natm,gto.PTR_COORD]
    ptr_coordL = numpy.vstack((ptr_coordL,ptr_coordL+1,ptr_coordL+2)).T.copy('C')
    for l, L1 in enumerate(Ls):
        env[ptr_coordL] = xyz + L1
        exp_Lk = numpy.einsum('k,ik->ik', expLk[l].conj(), expLk[:l+1])
        exp_Lk = numpy.asarray(exp_Lk, order='C')
        exp_Lk[l] = .5
        ints(exp_Lk, c_shls_slice)

    charge = cell.atom_charges()
    charge = numpy.append(charge, -charge)  # (charge-of-nuccell, charge-of-fakenuc)
    for k in range(nkpts):
        v = numpy.einsum('ijz,z->ij', buf[k], charge)
        buf[k] = lib.pack_tril(v + v.T.conj())

    if cell.dimension == 3:
        nucbar = sum([z/nuccell.bas_exp(i)[0] for i,z in enumerate(cell.atom_charges())])
        nucbar *= numpy.pi/cell.vol
        ovlp = cell.pbc_intor('cint1e_ovlp_sph', 1, lib.HERMITIAN, kpts)
        for k in range(nkpts):
            s = lib.pack_tril(ovlp[k])
            buf[k] += nucbar * s
    return buf

get_pp_loc_part1 = get_nuc

def get_pp(mydf, kpts=None):
    '''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed.
    '''
    cell = mydf.cell
    if kpts is None:
        kpts_lst = numpy.zeros((1,3))
    else:
        kpts_lst = numpy.reshape(kpts, (-1,3))
    nkpts = len(kpts_lst)

    vloc1 = mydf.get_nuc(kpts_lst)
    vloc2 = pseudo.pp_int.get_pp_loc_part2(cell, kpts_lst)
    vpp = pseudo.pp_int.get_pp_nl(cell, kpts_lst)
    for k in range(nkpts):
        vpp[k] += vloc1[k] + vloc2[k]

    if kpts is None or numpy.shape(kpts) == (3,):
        vpp = vpp[0]
    return vpp

class AFTDF(lib.StreamObject):
    '''Density expansion on plane waves
    '''
    def __init__(self, cell, kpts=numpy.zeros((1,3))):
        self.cell = cell
        self.stdout = cell.stdout
        self.verbose = cell.verbose
        self.max_memory = cell.max_memory
# For nuclear attraction integrals using Ewald-like technique.
# Set to 0 to use the regular reciprocal space Poisson-equation method.
        self.eta = estimate_eta(cell, cell.precision)

        self.kpts = kpts
        self.gs = cell.gs

        self.blockdim = 240 # to mimic molecular DF object

# Not input options
        self.exxdiv = None  # to mimic KRHF/KUHF object in function get_coulG
        self._keys = set(self.__dict__.keys())

    def dump_flags(self):
        log = logger.Logger(self.stdout, self.verbose)
        logger.info(self, '\n')
        logger.info(self, '******** %s flags ********', self.__class__)
        logger.info(self, 'gs = %s', self.gs)
        logger.info(self, 'len(kpts) = %d', len(self.kpts))
        logger.debug1(self, '    kpts = %s', self.kpts)

    def pw_loop(self, gs=None, kpti_kptj=None, q=None, shls_slice=None,
                max_memory=2000, aosym='s1', blksize=None):
        '''Plane wave part'''
        cell = self.cell
        if gs is None:
            gs = self.gs
        if kpti_kptj is None:
            kpti = kptj = numpy.zeros(3)
        else:
            kpti, kptj = kpti_kptj
        if q is None:
            q = kptj - kpti

        ao_loc = cell.ao_loc_nr()
        Gv, Gvbase, kws = cell.get_Gv_weights(gs)
        b = cell.reciprocal_vectors()
        gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
        ngs = gxyz.shape[0]

        if shls_slice is None:
            shls_slice = (0, cell.nbas, 0, cell.nbas)
        if aosym == 's2':
            assert(shls_slice[2] == 0)
            i0 = ao_loc[shls_slice[0]]
            i1 = ao_loc[shls_slice[1]]
            nij = i1*(i1+1)//2 - i0*(i0+1)//2
        else:
            ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
            nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
            nij = ni*nj

        if blksize is None:
            blksize = min(max(16, int(max_memory*1e6*.75/16/nij)), 16384)
            sublk = max(16, int(blksize//4))
        else:
            sublk = blksize
        buf = [numpy.zeros(nij*blksize, dtype=numpy.complex128)]
        pqkRbuf = numpy.empty(nij*sublk)
        pqkIbuf = numpy.empty(nij*sublk)

        if aosym == 's2':
            for p0, p1 in self.prange(0, ngs, blksize):
                aoao = ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym,
                                             b, gxyz[p0:p1], Gvbase, q,
                                             kptj.reshape(1,3), out=buf)[0]
                for i0, i1 in lib.prange(0, p1-p0, sublk):
                    nG = i1 - i0
                    pqkR = numpy.ndarray((nij,nG), buffer=pqkRbuf)
                    pqkI = numpy.ndarray((nij,nG), buffer=pqkIbuf)
                    pqkR[:] = aoao[i0:i1].real.T
                    pqkI[:] = aoao[i0:i1].imag.T
                    yield (pqkR, pqkI, p0+i0, p0+i1)
                aoao[:] = 0
        else:
            for p0, p1 in self.prange(0, ngs, blksize):
                #aoao = ft_ao.ft_aopair(cell, Gv[p0:p1], shls_slice, aosym,
                #                       b, Gvbase, gxyz[p0:p1], gs, (kpti, kptj), q)
                aoao = ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym,
                                             b, gxyz[p0:p1], Gvbase, q,
                                             kptj.reshape(1,3), out=buf)[0]
                for i0, i1 in lib.prange(0, p1-p0, sublk):
                    nG = i1 - i0
                    pqkR = numpy.ndarray((ni,nj,nG), buffer=pqkRbuf)
                    pqkI = numpy.ndarray((ni,nj,nG), buffer=pqkIbuf)
                    pqkR[:] = aoao[i0:i1].real.transpose(1,2,0)
                    pqkI[:] = aoao[i0:i1].imag.transpose(1,2,0)
                    yield (pqkR.reshape(-1,nG), pqkI.reshape(-1,nG), p0+i0, p0+i1)
                aoao[:] = 0

    def ft_loop(self, gs=None, q=numpy.zeros(3), kpts=None, shls_slice=None,
                max_memory=4000, aosym='s1'):
        '''
        Fourier transform iterator for all kpti which satisfy  2pi*N = (kpts - kpti - q)*a
        N = -1, 0, 1
        '''
        cell = self.cell
        if gs is None:
            gs = self.gs
        if kpts is None:
            assert(is_zero(q))
            kpts = self.kpts
        kpts = numpy.asarray(kpts)
        nkpts = len(kpts)

        ao_loc = cell.ao_loc_nr()
        b = cell.reciprocal_vectors()
        Gv, Gvbase, kws = cell.get_Gv_weights(gs)
        gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
        ngs = gxyz.shape[0]

        if shls_slice is None:
            shls_slice = (0, cell.nbas, 0, cell.nbas)
        if aosym == 's2':
            assert(shls_slice[2] == 0)
            i0 = ao_loc[shls_slice[0]]
            i1 = ao_loc[shls_slice[1]]
            nij = i1*(i1+1)//2 - i0*(i0+1)//2
        else:
            ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
            nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
            nij = ni*nj
        blksize = max(16, int(max_memory*.9e6/(nij*(nkpts+1)*16)))
        blksize = min(blksize, ngs, 16384)
        buf = [numpy.zeros(nij*blksize, dtype='D') for k in range(nkpts)]
        pqkRbuf = numpy.empty(nij*blksize)
        pqkIbuf = numpy.empty(nij*blksize)

        if aosym == 's2':
            for p0, p1 in self.prange(0, ngs, blksize):
                ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym,
                                      b, gxyz[p0:p1], Gvbase, q, kpts, out=buf)
                nG = p1 - p0
                for k in range(nkpts):
                    aoao = numpy.ndarray((nG,nij), dtype=numpy.complex128,
                                         order='F', buffer=buf[k])
                    pqkR = numpy.ndarray((nij,nG), buffer=pqkRbuf)
                    pqkI = numpy.ndarray((nij,nG), buffer=pqkIbuf)
                    pqkR[:] = aoao.real.T
                    pqkI[:] = aoao.imag.T
                    yield (k, pqkR, pqkI, p0, p1)
                    aoao[:] = 0  # == buf[k][:] = 0
        else:
            for p0, p1 in self.prange(0, ngs, blksize):
                ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym,
                                      b, gxyz[p0:p1], Gvbase, q, kpts, out=buf)
                nG = p1 - p0
                for k in range(nkpts):
                    aoao = numpy.ndarray((nG,ni,nj), dtype=numpy.complex128,
                                         order='F', buffer=buf[k])
                    pqkR = numpy.ndarray((ni,nj,nG), buffer=pqkRbuf)
                    pqkI = numpy.ndarray((ni,nj,nG), buffer=pqkIbuf)
                    pqkR[:] = aoao.real.transpose(1,2,0)
                    pqkI[:] = aoao.imag.transpose(1,2,0)
                    yield (k, pqkR.reshape(-1,nG), pqkI.reshape(-1,nG), p0, p1)
                    aoao[:] = 0  # == buf[k][:] = 0

    def prange(self, start, stop, step):
        return lib.prange(start, stop, step)

    def weighted_coulG(self, kpt=numpy.zeros(3), exx=False, gs=None):
        cell = self.cell
        if gs is None:
            gs = self.gs
        Gv, Gvbase, kws = cell.get_Gv_weights(gs)
        coulG = tools.get_coulG(cell, kpt, exx, self, gs, Gv)
        coulG *= kws
        return coulG

    _int_nuc_vloc = _int_nuc_vloc
    get_nuc = get_nuc
    get_pp = get_pp

    def get_jk(self, dm, hermi=1, kpts=None, kpts_band=None,
               with_j=True, with_k=True, exxdiv='ewald'):
        if kpts is None:
            if numpy.all(self.kpts == 0):
                # Gamma-point calculation by default
                kpts = numpy.zeros(3)
            else:
                kpts = self.kpts

        if kpts.shape == (3,):
            return aft_jk.get_jk(self, dm, hermi, kpts, kpts_band, with_j,
                                  with_k, exxdiv)

        vj = vk = None
        if with_k:
            vk = aft_jk.get_k_kpts(self, dm, hermi, kpts, kpts_band, exxdiv)
        if with_j:
            vj = aft_jk.get_j_kpts(self, dm, hermi, kpts, kpts_band)
        return vj, vk

    get_eri = get_ao_eri = aft_ao2mo.get_eri
    ao2mo = get_mo_eri = aft_ao2mo.general

    def update_mf(self, mf):
        mf = copy.copy(mf)
        mf.with_df = self
        return mf

################################################################################
# With this function to mimic the molecular DF.loop function, the pbc gamma
# point DF object can be used in the molecular code
    def loop(self):
        Lpq = None
        coulG = self.weighted_coulG()
        for pqkR, pqkI, p0, p1 in self.pw_loop(aosym='s2', blksize=self.blockdim):
            vG = numpy.sqrt(coulG[p0:p1])
            pqkR *= vG
            pqkI *= vG
            Lpq = lib.transpose(pqkR, out=Lpq)
            yield Lpq
            Lpq = lib.transpose(pqkI, out=Lpq)
            yield Lpq

    def get_naoaux(self):
        gs = numpy.asarray(mydf.gs)
        ngs = numpy.prod(gs*2+1)
        return ngs * 2

# Since the real-space lattice-sum for nuclear attraction is not implemented,
# use the 3c2e code with steep gaussians to mimic nuclear density
def _fake_nuc(cell):
    fakenuc = gto.Mole()
    fakenuc._atm = cell._atm.copy()
    fakenuc._atm[:,gto.PTR_COORD] = numpy.arange(gto.PTR_ENV_START,
                                                 gto.PTR_ENV_START+cell.natm*3,3)
    _bas = []
    _env = [0]*gto.PTR_ENV_START + [cell.atom_coords().ravel()]
    ptr = gto.PTR_ENV_START + cell.natm * 3
    half_sph_norm = .5/numpy.sqrt(numpy.pi)
    for ia in range(cell.natm):
        symb = cell.atom_symbol(ia)
        if symb in cell._pseudo:
            pp = cell._pseudo[symb]
            rloc, nexp, cexp = pp[1:3+1]
            eta = .5 / rloc**2
        else:
            eta = 1e16
        norm = half_sph_norm/gto.mole._gaussian_int(2, eta)
        _env.extend([eta, norm])
        _bas.append([ia, 0, 1, 1, 0, ptr, ptr+1, 0])
        ptr += 2
    fakenuc._bas = numpy.asarray(_bas, dtype=numpy.int32)
    fakenuc._env = numpy.asarray(numpy.hstack(_env), dtype=numpy.double)
    fakenuc.rcut = cell.rcut
    return fakenuc

if __name__ == '__main__':
    from pyscf.pbc import gto as pbcgto
    import pyscf.pbc.scf.hf as phf
    from pyscf.pbc import df
    cell = pbcgto.Cell()
    cell.verbose = 0
    cell.atom = 'C 0 0 0; C 1 1 1'
    cell.a = numpy.diag([4, 4, 4])
    cell.basis = 'gth-szv'
    cell.pseudo = 'gth-pade'
    cell.gs = [10, 10, 10]
    cell.build()
    k = numpy.ones(3)*.25
    v1 = AFTDF(cell).get_pp(k)
    print(abs(v1).sum(), 21.7504294462)