text:
#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import time
from pyscf import lib, __config__
from pyscf.grad import rhf as rhf_grad
from pyscf.soscf import ciah
import numpy as np
from scipy import linalg, optimize
from scipy.sparse import linalg as sparse_linalg
import time

default_level_shift = getattr(__config__, 'grad_lagrange_Gradients_level_shift', 1e-8)
default_conv_atol = getattr (__config__, 'grad_lagrange_Gradients_conv_atol', 1e-12)
default_conv_rtol = getattr (__config__, 'grad_lagrange_Gradients_conv_rtol', 1e-7)
default_max_cycle = getattr (__config__, 'grad_lagrange_Gradients_max_cycle', 50)

class Gradients (rhf_grad.GradientsMixin):
    r''' Dummy parent class for calculating analytical nuclear gradients using the technique of Lagrange multipliers:
    L = E + \sum_i z_i L_i
    dE/dx = \partial L/\partial x iff all L_i = 0 for the given wave function
    I.E., the Lagrange multipliers L_i cancel the direct dependence of the wave function on the nuclear coordinates
    and allow the Hellmann-Feynman theorem to be used for some non-variational methods. '''

    ################################## Child classes MUST overwrite the methods below ################################################

    def get_wfn_response (self, **kwargs):
        ''' Return first derivative of the energy wrt wave function parameters conjugate to the Lagrange multipliers.
            Used to calculate the value of the Lagrange multipliers. '''
        return np.zeros (nlag)

    def get_Aop_Adiag (self, **kwargs):
        ''' Return a function calculating Lvec . J_wfn, where J_wfn is the Jacobian of the Lagrange cofactors (e.g.,
            in state-averaged CASSCF, the Hessian of the state-averaged energy wrt wfn parameters) along with
            the diagonal of the Jacobian. '''
        def Aop (Lvec):
            return np.zeros (nlag)
        Adiag = np.zeros (nlag)
        return Aop, Adiag

    def get_ham_response (self, **kwargs):
        ''' Return expectation values <dH/dx> where x is nuclear displacement. I.E., the gradient if the method were variational. '''
        return np.zeros ((natm, 3))

    def get_LdotJnuc (self, Lvec, **kwargs):
        ''' Return Lvec . J_nuc, where J_nuc is the Jacobian of the Lagrange cofactors wrt nuclear displacement. This is the
            second term of the final gradient expectation value. '''
        return np.zeros ((natm, 3))

    ################################## Child classes SHOULD overwrite the methods below ##############################################

    def __init__(self, method, nlag):
        #self.mol = mol
        #self.base = method
        #self.verbose = mol.verbose
        #self.stdout = mol.stdout
        self.nlag = nlag
        #self.natm = mol.natm
        #self.atmlst = list (range (self.natm))
        #self.de = None
        #self._keys = set (self.__dict__.keys ())
        #--------------------------------------#
        self.level_shift = default_level_shift
        self.conv_atol = default_conv_atol
        self.conv_rtol = default_conv_rtol
        self.max_cycle = default_max_cycle
        rhf_grad.GradientsMixin.__init__(self, method)

    def debug_lagrange (self, Lvec, bvec, Aop, Adiag, **kwargs):
        lib.logger.debug (self, "{} gradient Lagrange factor debugging not enabled".format (self.base.__class__.__name__))
        pass

    def get_lagrange_callback (self, Lvec_last, itvec, geff_op):
        def my_call (x):
            itvec[0] += 1
            lib.logger.info (self, 'Lagrange optimization iteration {}, |geff| = {}, |dLvec| = {}'.format (itvec[0],
                linalg.norm (geff_op (x)), linalg.norm (x - Lvec_last)))
            Lvec_last[:] = x[:]
        return my_call

    def get_lagrange_precond (self, Adiag, level_shift=None, **kwargs):
        if level_shift is None: level_shift = self.level_shift
        return LagPrec (Adiag=Adiag, level_shift=level_shift, **kwargs)

    def get_init_guess (self, bvec, Adiag, Aop, precond):
        return precond (-bvec)

    ################################## Child classes SHOULD NOT overwrite the methods below ###########################################

    def solve_lagrange (self, Lvec_guess=None, level_shift=None, **kwargs):
        bvec = self.get_wfn_response (**kwargs)
        Aop, Adiag = self.get_Aop_Adiag (**kwargs)
        def my_geff (x):
            return bvec + Aop (x)
        Lvec_last = np.zeros_like (bvec)
        def my_Lvec_last ():
            return Lvec_last
        precond = self.get_lagrange_precond (Adiag, level_shift=level_shift, **kwargs)
        it = np.asarray ([0])
        lib.logger.debug (self, 'Lagrange multiplier determination intial gradient norm: {}'.format (linalg.norm (bvec)))
        my_call = self.get_lagrange_callback (Lvec_last, it, my_geff)
        Aop_obj = sparse_linalg.LinearOperator ((self.nlag,self.nlag), matvec=Aop, dtype=bvec.dtype)
        prec_obj = sparse_linalg.LinearOperator ((self.nlag,self.nlag), matvec=precond, dtype=bvec.dtype)
        x0_guess = self.get_init_guess (bvec, Adiag, Aop, precond)
        Lvec, info_int = sparse_linalg.cg (Aop_obj, -bvec, x0=x0_guess, tol=self.conv_rtol, atol=self.conv_atol, maxiter=self.max_cycle, callback=my_call, M=prec_obj)
        lib.logger.info (self, 'Lagrange multiplier determination {} after {} iterations\n   |geff| = {}, |Lvec| = {}'.format (
            ('converged','not converged')[bool (info_int)], it[0], linalg.norm (my_geff (Lvec)), linalg.norm (Lvec)))
        if info_int < 0: lib.logger.info (self, 'Lagrange multiplier determination error code {}'.format (info_int))
        return (info_int==0), Lvec, bvec, Aop, Adiag

    def kernel (self, level_shift=None, **kwargs):
        cput0 = (time.clock(), time.time())
        log = lib.logger.new_logger(self, self.verbose)
        if 'atmlst' in kwargs:
            self.atmlst = kwargs['atmlst']
        #self.natm = len (self.atmlst)

        if self.verbose >= lib.logger.WARN:
            self.check_sanity()
        if self.verbose >= lib.logger.INFO:
            self.dump_flags()

        conv, Lvec, bvec, Aop, Adiag = self.solve_lagrange (level_shift=level_shift, **kwargs)
        self.debug_lagrange (Lvec, bvec, Aop, Adiag, **kwargs)
        #if not conv: raise RuntimeError ('Lagrange multiplier determination not converged!')
        cput1 = lib.logger.timer (self, 'Lagrange gradient multiplier solution', *cput0)

        ham_response = self.get_ham_response (**kwargs)
        lib.logger.info(self, '--------------- %s gradient Hamiltonian response ---------------',
                    self.base.__class__.__name__)
        rhf_grad._write(self, self.mol, ham_response, self.atmlst)
        lib.logger.info(self, '----------------------------------------------')
        cput1 = lib.logger.timer (self, 'Lagrange gradient Hellmann-Feynman determination', *cput1)

        LdotJnuc = self.get_LdotJnuc (Lvec, **kwargs)
        lib.logger.info(self, '--------------- %s gradient Lagrange response ---------------',
                    self.base.__class__.__name__)
        rhf_grad._write(self, self.mol, LdotJnuc, self.atmlst)
        lib.logger.info(self, '----------------------------------------------')
        cput1 = lib.logger.timer (self, 'Lagrange gradient Jacobian', *cput1)

        self.de = ham_response + LdotJnuc
        log.timer('Lagrange gradients', *cput0)
        self._finalize()
        return self.de

    #def dump_flags(self):
    #    log = lib.logger.Logger(self.stdout, self.verbose)
    #    log.info('\n')
    #    if not self.base.converged:
    #        log.warn('Ground state method not converged')
    #    log.info('******** %s for %s ********',
    #             self.__class__, self.base.__class__)
    #    log.info('max_memory %d MB (current use %d MB)',
    #             self.max_memory, lib.current_memory()[0])
    #    return self

    #def _finalize (self):
    #    if self.verbose >= lib.logger.NOTE:
    #        lib.logger.note(self, '--------------- %s gradients ---------------',
    #                    self.base.__class__.__name__)
    #        rhf_grad._write(self, self.mol, self.de, self.atmlst)
    #        lib.logger.note(self, '----------------------------------------------')

class LagPrec (object):
    ''' A callable preconditioner for solving the Lagrange equations. Default is 1/(Adiagd+level_shift) '''

    def __init__(self, Adiag=None, level_shift=None, **kwargs):
        self.Adiag = Adiag
        self.level_shift = level_shift

    def __call__(self, x):
        Adiagd = self.Adiag + self.level_shift
        Adiagd[abs(Adiagd)<1e-8] = 1e-8
        x /= Adiagd
        return x


text:
import argparse
import collections
import glob
import os
import logging
import multiprocessing
import re
import shutil

from osgeo import gdal
import numpy
import ecoshard.geoprocessing
import taskgraph
import tempfile

import dnn_model
from utils.density_per_ha_to_total_per_pixel import \
    density_per_ha_to_total_per_pixel
from train_model import make_kernel_raster

gdal.SetCacheMax(2**27)

logging.basicConfig(
    level=logging.DEBUG,
    format=(
        '%(asctime)s (%(relativeCreated)d) %(levelname)s %(name)s'
        ' [%(funcName)s:%(lineno)d] %(message)s'))

LOGGER = logging.getLogger(__name__)
logging.getLogger('taskgraph').setLevel(logging.DEBUG)

# Working directories for substeps
def _mkdir(dir_path):
    """Safely make directory."""
    try:
        os.makedirs(dir_path)
    except OSError:
        pass
    return dir_path

WORKSPACE_DIR = _mkdir('./ipcc_vs_dnn_optimization')
CHURN_DIR = _mkdir(os.path.join(WORKSPACE_DIR, 'churn'))
BIOMASS_RASTER_DIR = _mkdir(
    os.path.join(WORKSPACE_DIR, 'biomass_rasters'))
MARGINAL_VALUE_WORKSPACE = _mkdir(
    os.path.join(WORKSPACE_DIR, 'marginal_value_rasters'))
OPTIMIZATION_WORKSPACE = _mkdir(
    os.path.join(WORKSPACE_DIR, 'optimization_workspaces'))
NEW_FOREST_MASK_DIR = _mkdir(
    os.path.join(WORKSPACE_DIR, 'new_forest_masks'))
IPCC_VS_DNN_DIR = _mkdir(
    os.path.join(WORKSPACE_DIR, 'ipcc_vs_dnn'))

MODEL_PATH = './models/model_400.dat'
BASE_DATA_DIR = 'workspace/ecoshard'

# *** DATA SECTION ***
# There are two landcover configurations, ESA and restoration of ESA
BASE_LULC_RASTER_PATH = os.path.join(
    BASE_DATA_DIR,
    'ESACCI-LC-L4-LCCS-Map-300m-P1Y-2014-v2.0.7_smooth_compressed.tif')
ESA_RESTORATION_SCENARIO_RASTER_PATH = os.path.join(
    BASE_DATA_DIR,
    'restoration_limited_md5_372bdfd9ffaf810b5f68ddeb4704f48f.tif')

# These are used in combination with an ESA landcover map to calculate carbon
CARBON_ZONES_VECTOR_PATH = os.path.join(
    BASE_DATA_DIR,
    'carbon_zones_md5_aa16830f64d1ef66ebdf2552fb8a9c0d.gpkg')
IPCC_CARBON_TABLE_PATH = os.path.join(
    BASE_DATA_DIR,
    'IPCC_carbon_table_md5_a91f7ade46871575861005764d85cfa7.csv')

# Constants useful for code readability
CARBON_MODEL_ID = os.path.basename(os.path.splitext(MODEL_PATH)[0])
IPCC_MODE = 'ipcc_mode'
# model mode is based off of carbon model ID
MODELED_MODE = f'modeled_mode_{CARBON_MODEL_ID}'
BASE_SCENARIO = 'base'
RESTORATION_SCENARIO = 'scenario'
FOREST_CODE = 50
TARGET_AREA_HA = 350000000*2
AREA_N_STEPS = 20
AREA_REPORT_STEP_LIST = numpy.linspace(
    TARGET_AREA_HA / AREA_N_STEPS, TARGET_AREA_HA, AREA_N_STEPS)
# number of pixels to blur to capture edge effect of marginal value
MARGINAL_VALUE_PIXEL_BLUR = 16

def _raw_basename(file_path):
    """Return just the filename without extension."""
    return os.path.basename(os.path.splitext(file_path)[0])[-20:]

def _sum_raster(raster_path):
    """Return sum of non-nodata values in ``raster_path``."""
    nodata = ecoshard.geoprocessing.get_raster_info(raster_path)['nodata'][0]
    running_sum = 0.0
    for _, raster_block in ecoshard.geoprocessing.iterblocks((raster_path, 1)):
        running_sum += numpy.sum(
            raster_block[~numpy.isclose(raster_block, nodata)])
    return running_sum

def _replace_value_by_mask(
        base_raster_path, replacement_value,
        replacement_mask_raster_path, target_replacement_raster_path):
    """Overwrite values in raster based on mask.

    Args:
        base_raster_path (str): base raster to modify
        replacement_value (numeric): value to write into base raster
            where the mask indicates.
        replacement_mask_raster_path (str): path to raster indicating (1) where
            a pixel should be replaced in base.
        target_replacement_raster_path (str): path to a target replacement
            raster.

    Returns:
        None
    """
    base_info = ecoshard.geoprocessing.get_raster_info(base_raster_path)
    ecoshard.geoprocessing.new_raster_from_base(
        base_raster_path, target_replacement_raster_path,
        base_info['datatype'], base_info['nodata'])
    target_raster = gdal.OpenEx(
        target_replacement_raster_path, gdal.OF_RASTER | gdal.GA_Update)
    target_band = target_raster.GetRasterBand(1)
    mask_raster = gdal.OpenEx(
        replacement_mask_raster_path, gdal.OF_RASTER | gdal.GA_Update)
    mask_band = mask_raster.GetRasterBand(1)

    for offset_dict, base_block in ecoshard.geoprocessing.iterblocks(
            (base_raster_path, 1)):
        mask_block = mask_band.ReadAsArray(**offset_dict)
        base_block[mask_block == 1] = replacement_value
        target_band.WriteArray(
            base_block, xoff=offset_dict['xoff'], yoff=offset_dict['yoff'])

    target_band = None
    target_raster = None

def _greedy_select_pixels_to_area(
        base_value_raster_path, workspace_dir, area_ha_to_step_report_list):
    """Greedy select pixels in base with a report every area steps.

    workspace_dir will contain a set of mask rasters with filenames of the form
    {area_selected}_mask_{base_id}.tif and a csv table with the filename
    {base_id}_{target_area_ha}_report.csv containing columns (area slected),
    (sum of value selected), (path to raster mask).

    Args:
        base_value_raster_path (str): path to raster with value pixels,
            preferably positive.
        workspace_dir (str): path to directory to write output files into.
        area_ha_to_step_report (list): list of areas in Ha to record.

    Returns:
        A tuple containing (path_to_taret_area_mask_raster,
            maximum area selected), where the raster is the largest amount
            selected and the value is the area that is selected, will either
            be very close to target_area_ha or the maximum available area.
    """
    raster_id = _raw_basename(base_value_raster_path)
    all_ones_raster_path = os.path.join(
        workspace_dir, f'all_ones_{raster_id}.tif')
    pixel_area_in_ha_raster_path = os.path.join(
        workspace_dir, f'pixel_area_in_ha_{raster_id}.tif')

    ecoshard.geoprocessing.new_raster_from_base(
        base_value_raster_path, all_ones_raster_path, gdal.GDT_Byte,
        [None], fill_value_list=[1])

    density_per_ha_to_total_per_pixel(
        all_ones_raster_path, 1.0,
        pixel_area_in_ha_raster_path)

    LOGGER.info(
        f'calculating greedy pixels for value raster {base_value_raster_path} '
        f'and area {pixel_area_in_ha_raster_path}')
    ecoshard.geoprocessing.greedy_pixel_pick_by_area_v2(
        (base_value_raster_path, 1), (pixel_area_in_ha_raster_path, 1),
        area_ha_to_step_report_list, workspace_dir)
    LOGGER.debug(
        f'done with greedy pixels for value raster {base_value_raster_path}')

def _create_marginal_value_layer(
        future_raster_path, base_raster_path,
        gaussian_blur_pixel_radius, mask_raster_path, target_raster_path):
    """Calculate marginal value layer.

    Calculated by taking the difference of future from base, Gaussian blurring
    that result by the given radius, and masking by the given raster mask.

    Args:
        future_raster_path (str): raster A, same nodata and size as B
        base_raster_path (str): raster B
        gaussian_blur_pixel_radius (int): number of pixels to blur out when
            determining marginal value of that pixel.
        mask_raster_path (str): path to raster where anything not 1 is masked
            to 0/nodata.
        target_diff_raster_path (str): result of A-B accounting for nodata.

    Returns:
        None
    """
    raster_info = ecoshard.geoprocessing.get_raster_info(future_raster_path)
    nodata = raster_info['nodata'][0]

    def _diff_op(a_array, b_array):
        """Return a-b and consider nodata."""
        result = numpy.copy(a_array)
        valid_mask = ~numpy.isclose(a_array, nodata)
        result[valid_mask] -= b_array[valid_mask]
        return result

    churn_dir = tempfile.mkdtemp(dir=os.path.dirname(target_raster_path))
    diff_raster_path = os.path.join(churn_dir, 'diff.tif')
    ecoshard.geoprocessing.raster_calculator(
        [(future_raster_path, 1), (base_raster_path, 1)], _diff_op,
        diff_raster_path, raster_info['datatype'], nodata)

    # Gaussian filter
    if gaussian_blur_pixel_radius is not None:
        kernel_raster_path = os.path.join(churn_dir, 'kernel.tif')
        mask_gf_path = os.path.join(churn_dir, 'gf.tif')
        if os.path.exists(mask_gf_path):
            os.remove(mask_gf_path)
        make_kernel_raster(
            gaussian_blur_pixel_radius, kernel_raster_path)
        ecoshard.geoprocessing.convolve_2d(
            (diff_raster_path, 1), (kernel_raster_path, 1), mask_gf_path,
            ignore_nodata_and_edges=False, mask_nodata=True,
            target_nodata=0.0)
    else:
        mask_gf_path = diff_raster_path

    def _mask_op(base_array, mask_array):
        """Return base where mask is 1, otherwise 0 or nodata."""
        result = numpy.copy(base_array)
        zero_mask = (~numpy.isclose(base_array, nodata)) & (mask_array != 1)
        result[zero_mask] = 0
        return result

    ecoshard.geoprocessing.raster_calculator(
        [(mask_gf_path, 1), (mask_raster_path, 1)], _mask_op,
        target_raster_path, raster_info['datatype'], nodata)

    shutil.rmtree(churn_dir)

def _diff_rasters(
        a_raster_path, b_raster_path, target_diff_raster_path):
    """Calculate a-b.

    Args:
        a_raster_path (str): raster A, same nodata and size as B
        b_raster_path (str): raster B
        target_diff_raster_path (str): result of A-B accounting for nodata.

    Returns:
        None
    """
    raster_info = ecoshard.geoprocessing.get_raster_info(a_raster_path)
    nodata = raster_info['nodata'][0]

    def _diff_op(a_array, b_array):
        """Return a-b and consider nodata."""
        result = numpy.copy(a_array)
        valid_mask = ~numpy.isclose(a_array, nodata)
        result[valid_mask] -= b_array[valid_mask]
        return result

    ecoshard.geoprocessing.raster_calculator(
        [(a_raster_path, 1), (b_raster_path, 1)], _diff_op,
        target_diff_raster_path, raster_info['datatype'], nodata)

def _calculate_new_forest(
        base_lulc_raster_path, future_lulc_raster_path,
        new_forest_mask_raster_path):
    """Calculate where there is new forest from base to future.

    Args:
        base_lulc_raster_path (str):
        future_lulc_raster_path (str):
        new_forest_mask_raster_path (str):

    Returns:
        None
    """
    FOREST_CODES = (50, 60, 61, 62, 70, 71, 72, 80, 81, 82, 90, 160, 170)

    def _mask_new_forest(base, future):
        """Remap values from ESA codes to basic MASK_TYPES."""
        result = numpy.empty(base.shape, dtype=numpy.uint8)
        base_forest = numpy.in1d(base, FOREST_CODES).reshape(result.shape)
        future_forest = numpy.in1d(future, FOREST_CODES).reshape(result.shape)
        result[:] = future_forest & ~base_forest
        return result

    ecoshard.geoprocessing.raster_calculator(
        [(base_lulc_raster_path, 1), (future_lulc_raster_path, 1)],
        _mask_new_forest, new_forest_mask_raster_path, gdal.GDT_Byte, None)

def _calculate_ipcc_biomass(
        landcover_raster_path, churn_dir, target_biomass_raster_path):
    """Calculate IPCC method for biomass for given landcover.

    Args:
        landcover_raster_path (str): path to ESA landcover raster.
        churn_dir (str): path to use for temporary files.
        target_biomass_raster_path (str): path to raster to create target
            biomass (not in density)

    Return:
        None
    """
    def _ipcc_carbon_op(
            lulc_array, zones_array, zone_lulc_to_carbon_map):
        """Map carbon to LULC/zone values and multiply by conversion map."""
        result = numpy.zeros(lulc_array.shape)
        for zone_id in numpy.unique(zones_array):
            if zone_id in zone_lulc_to_carbon_map:
                zone_mask = zones_array == zone_id
                result[zone_mask] = (
                    zone_lulc_to_carbon_map[zone_id][lulc_array[zone_mask]])
        return result

    def _parse_carbon_lulc_table(ipcc_carbon_table_path):
        """Parse out the IPCC carbon table by zone and lulc."""
        with open(IPCC_CARBON_TABLE_PATH, 'r') as carbon_table_file:
            header_line = carbon_table_file.readline()
            lulc_code_list = [
                int(lucode) for lucode in header_line.split(',')[1:]]
            max_code = max(lulc_code_list)

            zone_lucode_to_carbon_map = {}
            for line in carbon_table_file:
                split_line = line.split(',')
                if split_line[0] == '':
                    continue
                zone_id = int(split_line[0])
                zone_lucode_to_carbon_map[zone_id] = numpy.zeros(max_code+1)
                for lucode, carbon_value in zip(
                        lulc_code_list, split_line[1:]):
                    zone_lucode_to_carbon_map[zone_id][lucode] = float(
                        carbon_value)
        return zone_lucode_to_carbon_map

    rasterized_zones_raster_path = os.path.join(churn_dir, 'carbon_zones.tif')
    LOGGER.info(
        f'rasterize carbon zones of {landcover_raster_path} to '
        f'{rasterized_zones_raster_path}')
    ecoshard.geoprocessing.new_raster_from_base(
        landcover_raster_path, rasterized_zones_raster_path, gdal.GDT_Int32,
        [-1])
    ecoshard.geoprocessing.rasterize(
        CARBON_ZONES_VECTOR_PATH, rasterized_zones_raster_path,
        option_list=['ATTRIBUTE=CODE'])

    zone_lucode_to_carbon_map = _parse_carbon_lulc_table(
        IPCC_CARBON_TABLE_PATH)

    # uncommenting this so we get the density
    #biomass_per_ha_raster_path = os.path.join(churn_dir, 'biomass_per_ha.tif')
    ecoshard.geoprocessing.raster_calculator(
        [(landcover_raster_path, 1), (rasterized_zones_raster_path, 1),
         (zone_lucode_to_carbon_map, 'raw')],
        _ipcc_carbon_op, target_biomass_raster_path,
        gdal.GDT_Float32, -1)

    #density_per_ha_to_total_per_pixel(
    #    biomass_per_ha_raster_path, 1.0,
    #    target_biomass_raster_path)

def _calculate_modeled_biomass_from_mask(
        base_lulc_raster_path, new_forest_mask_raster_path,
        target_biomass_raster_path, task_graph):
    """Calculate new biomass raster from base layer and new forest mask.

    Args:
        base_lulc_raster_path (str): path to base ESA LULC raster.
        new_forest_mask_raster_path (str): path to raster that indicates
            where new forest is applied with a 1.
        target_biomass_raster_path (str): created by this function, a
            raster that has biomass per pixel for the scenario given by
            new_forest_mask_raster_path from base_lulc_raster_path.
        n_workers (int): number of workers to allow for reprojection.

    Returns:
        None
    """
    churn_dir = os.path.join(
        os.path.dirname(target_biomass_raster_path),
        os.path.basename(os.path.splitext(target_biomass_raster_path)[0]))

    # this raster is base with new forest in it
    converted_lulc_raster_path = os.path.join(churn_dir, 'converted_lulc.tif')
    LOGGER.info(
        f'creating converted LULC off of {base_lulc_raster_path} to '
        f'{converted_lulc_raster_path}')
    replace_value_by_mask_task = task_graph.add_task(
        func=_replace_value_by_mask,
        args=(
            base_lulc_raster_path, FOREST_CODE, new_forest_mask_raster_path,
            converted_lulc_raster_path),
        target_path_list=[converted_lulc_raster_path],
        task_name=f'replace by mask to {converted_lulc_raster_path}')
    replace_value_by_mask_task.join()

    # calculate biomass for that raster
    model_task = dnn_model.run_model(
        converted_lulc_raster_path,
        MODEL_PATH, target_biomass_raster_path, base_task_graph=task_graph)
    return model_task

def main():
    """Entry point."""
    parser = argparse.ArgumentParser(description='Build marginal value map')
    parser.add_argument(
        '--bounding_box', type=float, nargs=4, help='xmin, xmax, ymin, ymax')
    args = parser.parse_args()

    if args.bounding_box:
        bb_string = f"_{'_'.join([f'{v:.2f}' for v in args.bounding_box])}"
    else:
        bb_string = ''

    task_graph = taskgraph.TaskGraph(
        WORKSPACE_DIR, multiprocessing.cpu_count())

    dnn_model.download_data(task_graph, dnn_model.BOUNDING_BOX)

    unique_scenario_id = f'''{
        _raw_basename(BASE_LULC_RASTER_PATH)}_{
        _raw_basename(ESA_RESTORATION_SCENARIO_RASTER_PATH)}{bb_string}'''

    LOGGER.info(f'calculate new forest mask on {BASE_LULC_RASTER_PATH}')
    new_forest_raster_path = os.path.join(
        NEW_FOREST_MASK_DIR, f'{unique_scenario_id}.tif')
    new_forest_mask_task = task_graph.add_task(
        func=_calculate_new_forest,
        args=(
            BASE_LULC_RASTER_PATH, ESA_RESTORATION_SCENARIO_RASTER_PATH,
            new_forest_raster_path),
        target_path_list=[new_forest_raster_path],
        task_name=f'create forest mask for {new_forest_raster_path}')

    # modeled_biomass_raster_task_dict indexed by
    #   [MODELED_MODE/IPCC_MODE] -> [BASE_SCENARIO/RESTORATION_SCENARIO]
    modeled_biomass_raster_task_dict = collections.defaultdict(dict)
    for scenario_id, landcover_raster_path in [
            (BASE_SCENARIO, BASE_LULC_RASTER_PATH),
            (RESTORATION_SCENARIO, ESA_RESTORATION_SCENARIO_RASTER_PATH)]:
        # create churn directory and id for modeled biomass.
        base_landcover_id = os.path.basename(
            os.path.splitext(landcover_raster_path)[0])

        # calculated modeled biomass
        LOGGER.info(
            f'model biomass {MODELED_MODE} for {base_landcover_id}/'
            f'{scenario_id}')
        modeled_biomass_raster_path = os.path.join(
            BIOMASS_RASTER_DIR,
            f'biomass_{MODELED_MODE}_{scenario_id}.tif')
        biomass_model_task = dnn_model.run_model(
            landcover_raster_path, MODEL_PATH,
            modeled_biomass_raster_path, base_task_graph=task_graph)

        modeled_biomass_raster_task_dict[MODELED_MODE][scenario_id] = \
            (modeled_biomass_raster_path, biomass_model_task)

        # calculate IPCC biomass
        LOGGER.info(
            f'calculate IPCC method for {base_landcover_id}/'
            f'{scenario_id}')
        target_ipcc_biomass_path = os.path.join(
            BIOMASS_RASTER_DIR,
            f'biomass_per_ha_{IPCC_MODE}_{scenario_id}.tif')
        ipcc_churn_dir = os.path.join(
            CHURN_DIR, f'churn_{base_landcover_id}_{IPCC_MODE}')
        biomass_ipcc_task = task_graph.add_task(
            func=_calculate_ipcc_biomass,
            args=(
                landcover_raster_path, ipcc_churn_dir,
                target_ipcc_biomass_path),
            target_path_list=[target_ipcc_biomass_path],
            task_name=f'calculate biomass{IPCC_MODE} for {scenario_id}')
        modeled_biomass_raster_task_dict[IPCC_MODE][scenario_id] = \
            (target_ipcc_biomass_path, biomass_ipcc_task)

    LOGGER.info('create marginal value maps')
    # this will have (mode, task, dir) tuples for this section
    optimization_mode_task_dir_list = []
    for model_mode in [MODELED_MODE, IPCC_MODE]:
        marginal_value_biomass_raster = os.path.join(
            MARGINAL_VALUE_WORKSPACE,
            f'marginal_value_biomass_{model_mode}.tif')
        restoration_biomass_raster, restoration_task = \
            modeled_biomass_raster_task_dict[model_mode][RESTORATION_SCENARIO]
        base_biomass_raster, base_task = \
            modeled_biomass_raster_task_dict[model_mode][BASE_SCENARIO]

        marginal_value_task = task_graph.add_task(
            func=_create_marginal_value_layer,
            args=(
                restoration_biomass_raster,
                base_biomass_raster,
                (MARGINAL_VALUE_PIXEL_BLUR
                 if model_mode == MODELED_MODE else None),
                new_forest_raster_path,
                marginal_value_biomass_raster),
            target_path_list=[marginal_value_biomass_raster],
            dependent_task_list=[
                restoration_task, base_task, new_forest_mask_task],
            task_name=(
                f'''calc marginal value for {restoration_biomass_raster} '''
                f'''and {base_biomass_raster}'''))

        LOGGER.info(
            f'create optimal land selection mask to target '
            f'{TARGET_AREA_HA} ha')
        optimization_dir = _mkdir(os.path.join(
            OPTIMIZATION_WORKSPACE,
            f'optimization_{unique_scenario_id}_{model_mode}'))
        # returns a (optimal mask, area selected) tuple
        optimization_task = task_graph.add_task(
            func=_greedy_select_pixels_to_area,
            args=(
                marginal_value_biomass_raster, optimization_dir,
                AREA_REPORT_STEP_LIST),
            target_path_list=[os.path.join(optimization_dir, f'marginal_value_biomass_{model_mode}_greedy_pick.csv')],
            dependent_task_list=[marginal_value_task],
            task_name=f'optimize on {marginal_value_biomass_raster}')
        optimization_mode_task_dir_list.append(
            (model_mode, optimization_task, optimization_dir))

        task_graph.join()
        task_graph.close()
        return

    # indexed by MODELED_MODE vs. IPCC_MODE then by area of the new forest
    optimization_biomass_area_path_task_dict = \
        collections.defaultdict(dict)
    for (model_mode, optimization_task, optimization_dir) in \
            optimization_mode_task_dir_list:
        # okay to join here because it's going to trigger a whole set of
        # other tasks and nothing can be done until this one is ready anyway
        optimization_task.join()

        optimization_biomass_dir = _mkdir(os.path.join(
            OPTIMIZATION_WORKSPACE,
            f'biomass_{unique_scenario_id}_{model_mode}'))

        for optimal_mask_raster_path in glob.glob(
                os.path.join(optimization_dir, 'mask_*.tif')):
            optimization_biomass_raster_path = os.path.join(
                optimization_biomass_dir,
                f'''biomass_per_pixel_{
                    _raw_basename(optimal_mask_raster_path)}.tif''')
            model_task = _calculate_modeled_biomass_from_mask(
                BASE_LULC_RASTER_PATH, optimal_mask_raster_path,
                optimization_biomass_raster_path, task_graph)
            mask_area = float(re.match(
                r'mask_(.*)\.tif', os.path.basename(
                    optimal_mask_raster_path)).group(1))
            optimization_biomass_area_path_task_dict[
                model_mode][mask_area] = (
                    optimization_biomass_raster_path, model_task)

    # TODO: calculate difference between modeled vs IPCC
    task_graph.join()
    LOGGER.info(
        'calculate difference between modeled biomass optimization and IPCC '
        'optimization')
    mask_areas = sorted([
        float(x) for x in
        optimization_biomass_area_path_task_dict[MODELED_MODE].keys()])
    modeled_diff_ipcc_biomass_sum_task_list = []
    modeled_diff_mode_base_biomass_sum_task_list = \
        collections.defaultdict(list)
    LOGGER.debug(f'********* {mask_areas} {optimization_dir}')
    for mask_area in mask_areas:
        model_biomass_raster_path, dnn_model_task = \
            optimization_biomass_area_path_task_dict[MODELED_MODE][mask_area]
        ipcc_biomass_raster_path, ipcc_model_task = \
            optimization_biomass_area_path_task_dict[IPCC_MODE][mask_area]
        modeled_vs_ipcc_optimal_biomass_diff_raster_path = os.path.join(
            IPCC_VS_DNN_DIR,
            f'modeled_vs_ipcc_diff_{mask_area}_ha.tif')
        diff_task = task_graph.add_task(
            func=_diff_rasters,
            args=(
                model_biomass_raster_path,
                ipcc_biomass_raster_path,
                modeled_vs_ipcc_optimal_biomass_diff_raster_path),
            target_path_list=[
                modeled_vs_ipcc_optimal_biomass_diff_raster_path],
            dependent_task_list=[dnn_model_task, ipcc_model_task],
            task_name=f'''modeled diff ipcc {
                modeled_vs_ipcc_optimal_biomass_diff_raster_path}''')

        sum_task = task_graph.add_task(
            func=_sum_raster,
            args=(modeled_vs_ipcc_optimal_biomass_diff_raster_path,),
            store_result=True,
            dependent_task_list=[diff_task],
            task_name=f'''sum the modeled vs. ippc diff for {
                modeled_vs_ipcc_optimal_biomass_diff_raster_path}''')
        modeled_diff_ipcc_biomass_sum_task_list.append(sum_task)

        LOGGER.info(
            'subtract modeled and ipcc from base modeled to get the gain')
        modeled_vs_base_biomass_diff_raster_path = os.path.join(
            IPCC_VS_DNN_DIR, f'modeled_vs_base_{mask_area}_ha.tif')
        ipcc_vs_base_biomass_diff_raster_path = os.path.join(
            IPCC_VS_DNN_DIR, f'ipcc_vs_base_{mask_area}_ha.tif')
        modeled_base_biomass_raster_path = (
            modeled_biomass_raster_task_dict[MODELED_MODE][BASE_SCENARIO][0])
        for modeled_path, target_diff_path, mode, model_task in [
                (model_biomass_raster_path,
                 modeled_vs_base_biomass_diff_raster_path, MODELED_MODE, dnn_model_task),
                (ipcc_biomass_raster_path,
                 ipcc_vs_base_biomass_diff_raster_path, IPCC_MODE, ipcc_model_task)]:
            diff_task = task_graph.add_task(
                func=_diff_rasters,
                args=(
                    modeled_path, modeled_base_biomass_raster_path,
                    target_diff_path),
                dependent_task_list=[model_task],
                target_path_list=[target_diff_path],
                task_name=f'modeled diff ipcc {target_diff_path}')

            sum_task = task_graph.add_task(
                func=_sum_raster,
                args=(target_diff_path,),
                store_result=True,
                dependent_task_list=[diff_task],
                task_name=f'''sum the modeled/ipcc vs. base for {
                    target_diff_path}''')
            modeled_diff_mode_base_biomass_sum_task_list[mode].append(sum_task)

    task_graph.join()

    LOGGER.info('report')
    report_csv_path = os.path.join(
        WORKSPACE_DIR, f'''report_{_raw_basename(BASE_LULC_RASTER_PATH)}_{
        _raw_basename(ESA_RESTORATION_SCENARIO_RASTER_PATH)}.csv''')
    with open(report_csv_path, 'w') as report_csv_file:
        report_csv_file.write(
            'area,biomass gain modeled driven, biomass gain IPCC driven, '
            'modeled vs ipcc diff\n')

        for (area, biomass_modeled_gain, ipcc_modeled_gain,
                modeled_vs_ipcc) in zip(
                mask_areas,
                modeled_diff_mode_base_biomass_sum_task_list[MODELED_MODE],
                modeled_diff_mode_base_biomass_sum_task_list[IPCC_MODE],
                modeled_diff_ipcc_biomass_sum_task_list):
            report_csv_file.write(
                f'{area},{biomass_modeled_gain.get()},'
                f'{ipcc_modeled_gain.get()},{modeled_vs_ipcc.get()}\n')
    task_graph.join()
    task_graph.close()

if __name__ == '__main__':
    main()

ext: py

reason: import numpy


text:
import networkx as nx
import numpy as np
import os
import time
import pandas as pd
from sscd_s import SSCDS
from ssnmf import SSNMF
import scipy.sparse as ssp

"""Functions definitions"""
def generate_synthetic_graph(n_node_par_com, n_com=2, p_in=0.1, p_out=0.001):
    sizes = [n_node_par_com] * n_com
    g = nx.random_partition_graph(sizes, p_in, p_out)
    edges = g.edges()
    return edges

def make_matrix_from_elist(elist, nn):
    mat = ssp.lil_matrix((nn,nn))
    ind = np.array(elist)
    mat[ind[:,0],ind[:,1]] = 1
    return mat.toarray()

def calculate_loss(edge_list, W, H, mlambda, const):
    nn = W.shape[0]
    X = make_matrix_from_elist(edge_list, nn)
    O = make_matrix_from_elist(const, nn)
    lse = np.sum(np.power(X - np.matmul(W,H.T),2))
    reg = np.matmul(H.T, np.matmul(O, H)).trace()
    loss = lse + mlambda * reg
    return loss

def export_cost(path, cost_list):
    with open(path, "w") as f:
        for c in cost_list:
            f.write(str(c) + "\n")

mlambda = 1.0
threads = 8
lr_adam = 0.1
K = 10
abs_adam = SSCDS(K, method="adam", positivate="abs", mlambda=mlambda, learning_rate=lr_adam, threads=threads)

model = abs_adam

start =- time.time()
dblp = pd.read_pickle("data/dblp_edge.pkl")
const = [(1,2)]
W, H, mean_cost, cost_list = model.fit_and_transform(dblp, const)
print("Takes " + str(time.time()-start) + "[s]")

max_stars_repo_name: nukui-s/sscomdetection

max_issues_repo_name: nukui-s/sscomdetection

max_forks_repo_name: nukui-s/sscomdetection

reason: import numpy,import scipy,import networkx


text:
import pylab as p
import numpy as n
import cosmolopy.perturbation as pb
import cosmolopy.density as cd
from scipy.integrate import quad
from scipy.interpolate import interp1d, interp2d
from scipy.interpolate import RectBivariateSpline as RBS

def del0over1plusz(th):
	return 3./20*(6*(th-n.sin(th)))*(2./3)
def redshift(del0,th):
	RHS = del0over1plusz(th)/del0
	return 1/RHS-1

del0 = 3.
theta = n.arange(1.2,4.,0.01)
#Z = redshift(del0,theta)
d0over1plusz = del0over1plusz(theta)

max_stars_repo_name: yunfanz/ReionBub

max_issues_repo_name: yunfanz/ReionBub

max_forks_repo_name: yunfanz/ReionBub

reason: import numpy,from scipy


text:
import matplotlib.pyplot as plt
import numpy as np
import sys
from gym_unity.envs import UnityEnv
from mlagents.envs import UnityEnvironment

SIZE_OBSERVATION = 112

class SocTwoEnv():
    """
        parameters:
            env_path: the binary file which is built by Unity.
            train_mode: It's True when u want to train the brain.
            n_str: number of agent of striker in the scene.
            n_goalie: number of agent of goalie in the scene.

        ***********************************************************

        Store "action" of each agent in act_str_hist and act_goalie_hist
        respectively.

        Store "Observation" or "State" of each agent in observation_str_hist
        and observation_goalie_hist respectively.

    """
    def __init__(self,
                 env_path,
                 worker_id,
                 train_mode=True,
                 n_str=16,
                 n_goalie=16):
        self.env = UnityEnvironment(file_name=env_path, worker_id=0)
        self.striker_brain_name, self.goalie_brain_name = self.env.brain_names
        self.striker_brain = self.env.brains[self.striker_brain_name]
        self.goalie_brain = self.env.brains[self.goalie_brain_name]
        self.done_str = [False] * 16
        self.done_goalie = [False] * 16
        self.train_mode = train_mode
        self.done_hist_str = [False] * 16
        self.done_hist_goalie = [False] * 16
        self.episode_str_rewards = 0
        self.episode_goalie_rewards = 0
        self.n_str = n_str
        self.n_goalie = n_goalie
        self.act_str_hist = [[] for x in range(n_str)]
        self.act_goalie_hist = [[] for x in range(n_goalie)]
        self.observation_str_hist = [[] for x in range(SIZE_OBSERVATION)]
        self.observation_goalie_hist = [[] for x in range(SIZE_OBSERVATION)]
        self.observation_str = None
        self.observation_goalie = None
        return

    def reset(self):
        """
            Reset the all environments and agents.
        """
        self.env_info_str = self.env.reset(
            train_mode=self.train_mode)[self.striker_brain_name]
        print("env_info_str", self.env_info_str)
        self.env_info_goalie = self.env.reset(
            train_mode=self.train_mode)[self.goalie_brain_name]
        self.episode_rewards = 0
        self.done_str = [False] * 16
        self.done_goalie = [False] * 16
        self.done_hist_str = np.array([False] * 16)
        self.done_hist_goalie = np.array([False] * 16)
        return {'str': self.env_info_str, 'goalie': self.env_info_goalie}

    def step(self, action_str, action_goalie):
        """
            In each timestep, give each striker and goalie a instruction
            to do action. And then, get the current observation stored
            at observation_str and observation_goalie.
        """
        self.env_info = self.env.step({
            self.striker_brain_name: action_str,
            self.goalie_brain_name: action_goalie
        })
        self.observation_str = np.array(
            self.env_info[self.striker_brain_name].vector_observations)
        self.observation_goalie = np.array(
            self.env_info[self.goalie_brain_name].vector_observations)
        return self.env_info

    def reward(self):
        """
            return the rewards of striker and goalie respectively.
        """
        self.episode_str_rewards = np.array(
            self.env_info[self.striker_brain_name].rewards)
        self.episode_goalie_rewards = np.array(
            self.env_info[self.goalie_brain_name].rewards)
        return self.episode_str_rewards, self.episode_goalie_rewards

    def close(self):
        """
            Close the simulation Unity environment.
        """
        self.env.close()

    def done(self):
        """
            Check whether each agent has finished at their own episode,
            true means it has finished.
        """
        self.done_str = np.array(
            self.env_info[self.striker_brain_name].local_done)
        self.done_goalie = np.array(
            self.env_info[self.goalie_brain_name].local_done)

    def reset_some_agents(self, str_arg, goalie_arg):
        """
            params:
                str_arg, mark which striker's history that wants to be cleared.
                goalie_arg, mark which goalie's history that wants to be cleared.

            Clear the history of specific agents.

        """
        for i in str_arg:
            self.act_str_hist[i[0]] = []
            self.observation_str_hist[i[0]] = []
        for i in goalie_arg:
            self.act_goalie_hist[i[0]] = []
            self.observation_goalie_hist[i[0]] = []

if __name__ == "__main__":

    env_Path = r'.\env\windows\soccer_easy\Unity Environment.exe'
    soc_env = SocTwoEnv(env_Path, worker_id=0, train_mode=True)
    soc_env.reset()  # Don't touch me!
    episode = 0
    while episode < 10:
        action_size_str = soc_env.striker_brain.vector_action_space_size
        action_size_goalie = soc_env.goalie_brain.vector_action_space_size

        # randomly generate some actions for each agent.
        action_str = np.column_stack([
            np.random.randint(0, 7, size=16)
            for i in range(len(action_size_str))
        ])
        action_goalie = np.column_stack([
            np.random.randint(0, action_size_goalie[i], size=16)
            for i in range(len(action_size_goalie))
        ])

        # store the action of agent in list
        for i in range(soc_env.n_goalie):
            (soc_env.act_goalie_hist[i]).append(action_goalie[i])
        for i in range(soc_env.n_str):
            (soc_env.act_str_hist[i]).append(action_str[i])
        soc_env.step(action_str, action_goalie)
        # store the observation(state) of agent in list
        for i in range(soc_env.n_goalie):
            (soc_env.observation_goalie_hist[i]).append(
                soc_env.observation_goalie[i])
        for i in range(soc_env.n_str):
            (soc_env.observation_str_hist[i]).append(
                soc_env.observation_str[i])

        soc_env.done()
        if True in soc_env.done_goalie:
            soc_env.reward()
            print("episode: ", episode, "*" * 10)
            arg_done_goalie = np.argwhere(soc_env.done_goalie == True)
            for i in arg_done_goalie:
                # print("goalie %d"%(i[0]))
                # print("action", soc_env.act_goalie_hist[i[0]])
                # print("Observation", soc_env.observation_goalie_hist[i[0]])
                # print("reword", soc_env.episode_goalie_rewards[i][0])
                pass

            arg_done_str = np.argwhere(soc_env.done_str == True)
            for i in arg_done_str:
                # print("str %d"%(i[0]))
                # print("action", soc_env.act_str_hist[i[0]])
                # print("Observation", soc_env.observation_str_hist[i[0]])
                # print("reword", soc_env.episode_str_rewards[i][0])
                pass
            soc_env.reset_some_agents(arg_done_str, arg_done_goalie)
            print("*" * 25)
            episode += 1
            c = 0
    soc_env.close()


text:
from sympy import *

# commutative=False 用来防止x.g1变成g1.x
x, y, a, b, c, d, g1, g2, g11, g21 = symbols("x, y, a, b, c, d, g1, g2, g1',g2'", commutative=False)
E, F, G = symbols("E, F, G", commutative=False)
GG = Matrix([[a, c], [b, d]])
v = Matrix([UnevaluatedExpr(x), UnevaluatedExpr(y)])
v_alg = UnevaluatedExpr(x)*UnevaluatedExpr(g1)+UnevaluatedExpr(y)*UnevaluatedExpr(g2)
pprint(G)
pprint(v)
pprint(v.T)

exp = UnevaluatedExpr(v.T) * UnevaluatedExpr(G) * UnevaluatedExpr(v)
pprint(exp)
pprint(exp.doit())
print("---------------")
vv = UnevaluatedExpr(v.T*v)
P = UnevaluatedExpr(x)*UnevaluatedExpr(g1) + UnevaluatedExpr(y)*UnevaluatedExpr(g2)
PP = P*P
print_latex(PP.expand())
GG = Matrix([[E, F], [F, G]])
contract_matrix = UnevaluatedExpr(v.T)*UnevaluatedExpr(GG)*UnevaluatedExpr(v)
pprint(contract_matrix)
print_latex(contract_matrix)
# pprint(vv)
# vv_alg = v_alg * v_alg
# pprint(vv_alg.expand())

max_stars_repo_name: baijianhua/pymath

max_issues_repo_name: baijianhua/pymath

max_forks_repo_name: baijianhua/pymath

reason: from sympy


text:
# *****************************************************************************
#
# Copyright (c) 2019, the Perspective Authors.
#
# This file is part of the Perspective library, distributed under the terms of
# the Apache License 2.0.  The full license can be found in the LICENSE file.
#
import six
import os
import numpy as np
import pandas as pd
import pytz
from io import StringIO
from datetime import date, datetime
from perspective.table import Table
from pytest import mark
IS_WIN = os.name == 'nt'

class TestToFormat(object):

    # to_records

    def test_to_records_int(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records() == data

    def test_to_records_float(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records() == data

    def test_to_records_string(self):
        data = [{"a": "string1", "b": "string2"}, {"a": "string3", "b": "string4"}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records() == data

    def test_to_records_date(self):
        today = date.today()
        dt = datetime(today.year, today.month, today.day)
        data = [{"a": today, "b": "string2"}, {"a": today, "b": "string4"}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records() == [{"a": dt, "b": "string2"}, {"a": dt, "b": "string4"}]

    def test_to_records_date_no_dst(self):
        # make sure that DST does not affect the way we read dates - if tm_dst in `t_date::get_tm()` isn't set to -1, it could reverse 1hr by assuming DST is not in effect.
        today = date.today()
        dt = datetime(today.year, today.month, today.day)
        data = [{"a": today, "b": "string2"}, {"a": today, "b": "string4"}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records() == [{"a": dt, "b": "string2"}, {"a": dt, "b": "string4"}]

    def test_to_records_date_str(self):
        data = [{"a": "03/11/2019", "b": "string2"}, {"a": "03/12/2019", "b": "string4"}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records() == [{"a": datetime(2019, 3, 11), "b": "string2"}, {"a": datetime(2019, 3, 12), "b": "string4"}]

    def test_to_records_date_str_month_first(self):
        data = [{"a": "1/2/2019", "b": "string2"}, {"a": "3/4/2019", "b": "string4"}]
        tbl = Table(data)
        view = tbl.view()
        assert view.schema() == {"a": date, "b": str}
        assert view.to_records() == [{"a": datetime(2019, 1, 2), "b": "string2"}, {"a": datetime(2019, 3, 4), "b": "string4"}]

    def test_to_records_date_str_month_ymd(self):
        data = [{"a": "2019/01/02", "b": "string2"}, {"a": "2019/03/04", "b": "string4"}]
        tbl = Table(data)
        view = tbl.view()
        assert view.schema() == {"a": date, "b": str}
        assert view.to_records() == [{"a": datetime(2019, 1, 2), "b": "string2"}, {"a": datetime(2019, 3, 4), "b": "string4"}]

    def test_to_records_datetime(self):
        dt = datetime(2019, 9, 10, 19, 30, 59, 515000)
        data = [{"a": dt, "b": "string2"}, {"a": dt, "b": "string4"}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records() == data  # should have symmetric input/output

    def test_to_records_datetime_str(self):
        data = [{"a": "03/11/2019 3:15PM", "b": "string2"}, {"a": "3/11/2019 3:20PM", "b": "string4"}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records() == [{"a": datetime(2019, 3, 11, 15, 15), "b": "string2"}, {"a": datetime(2019, 3, 11, 15, 20), "b": "string4"}]

    def test_to_records_datetime_str_tz(self):
        dt = "2019/07/25T15:30:00+00:00"
        data = [{"a": dt}, {"a": dt}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records()
        for r in records:
            r["a"] = r["a"].replace(tzinfo=pytz.utc)
        assert records == [{"a": datetime(2019, 7, 25, 15, 30, tzinfo=pytz.utc)}, {"a": datetime(2019, 7, 25, 15, 30, tzinfo=pytz.utc)}]

    def test_to_records_datetime_ms_str(self):
        data = [{"a": "03/11/2019 3:15:15.999PM"}, {"a": "3/11/2019 3:15:16.001PM"}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records() == [{"a": datetime(2019, 3, 11, 15, 15, 15, 999000)}, {"a": datetime(2019, 3, 11, 15, 15, 16, 1000)}]

    def test_to_records_none(self):
        data = [{"a": None, "b": 1}, {"a": None, "b": 2}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records() == data

    def test_to_records_one(self):
        data = [{"a": 1, "b": "string1"}, {"a": 1, "b": "string2"}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"]
        )
        assert view.to_records() == [
            {"__ROW_PATH__": [], "a": 2, "b": 2}, {"__ROW_PATH__": [1], "a": 2, "b": 2}
        ]

    def test_to_records_two(self):
        data = [{"a": 1, "b": "string1"}, {"a": 1, "b": "string2"}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"]
        )
        assert view.to_records() == [
            {"__ROW_PATH__": [], "string1|a": 1, "string1|b": 1, "string2|a": 1, "string2|b": 1},
            {"__ROW_PATH__": [1], "string1|a": 1, "string1|b": 1, "string2|a": 1, "string2|b": 1},
        ]

    def test_to_records_column_only(self):
        data = [{"a": 1, "b": "string1"}, {"a": 1, "b": "string2"}]
        tbl = Table(data)
        view = tbl.view(
            column_pivots=["b"]
        )
        assert view.to_records() == [
            {"string1|a": 1, "string1|b": "string1", "string2|a": None, "string2|b": None},
            {"string1|a": None, "string1|b": None, "string2|a": 1, "string2|b": "string2"},
        ]

    # to_dict

    def test_to_dict_int(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_dict() == {
            "a": [1, 3],
            "b": [2, 4]
        }

    def test_to_dict_float(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_dict() == {
            "a": [1.5, 3.5],
            "b": [2.5, 4.5]
        }

    def test_to_dict_date(self):
        today = date.today()
        dt = datetime(today.year, today.month, today.day)
        data = [{"a": today, "b": 2}, {"a": today, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_dict() == {
            "a": [dt, dt],
            "b": [2, 4]
        }

    def test_to_dict_datetime(self):
        dt = datetime(2019, 3, 15, 20, 30, 59, 6000)
        data = [{"a": dt, "b": 2}, {"a": dt, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_dict() == {
            "a": [dt, dt],
            "b": [2, 4]
        }

    def test_to_dict_bool(self):
        data = [{"a": True, "b": False}, {"a": True, "b": False}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_dict() == {
            "a": [True, True],
            "b": [False, False]
        }

    def test_to_dict_string(self):
        data = [{"a": "string1", "b": "string2"}, {"a": "string3", "b": "string4"}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_dict() == {
            "a": ["string1", "string3"],
            "b": ["string2", "string4"]
        }

    def test_to_dict_none(self):
        data = [{"a": None, "b": None}, {"a": None, "b": None}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_dict() == {
            "a": [None, None],
            "b": [None, None]
        }

    def test_to_dict_one(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"]
        )
        assert view.to_dict() == {
            "__ROW_PATH__": [[], [1]],
            "a": [2, 2],
            "b": [4, 4]
        }

    def test_to_dict_two(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"]
        )
        assert view.to_dict() == {
            "__ROW_PATH__": [[], [1]],
            "2|a": [2, 2],
            "2|b": [4, 4]
        }

    def test_to_dict_column_only(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            column_pivots=["b"]
        )
        assert view.to_dict() == {
            "2|a": [1, 1],
            "2|b": [2, 2],
        }

    def test_to_dict_one_no_columns(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            columns=[]
        )
        assert view.to_dict() == {"__ROW_PATH__": [[], [1]]}

    def test_to_dict_two_no_columns(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"],
            columns=[]
        )
        assert view.to_dict() == {
            "__ROW_PATH__": [[], [1]]
        }

    def test_to_dict_column_only_no_columns(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            column_pivots=["b"],
            columns=[]
        )
        assert view.to_dict() == {}

    # to_numpy

    def test_to_numpy_int(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        v = view.to_numpy()
        assert np.array_equal(v["a"], np.array([1, 3]))
        assert np.array_equal(v["b"], np.array([2, 4]))

    def test_to_numpy_float(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        v = view.to_numpy()
        assert np.array_equal(v["a"], np.array([1.5, 3.5]))
        assert np.array_equal(v["b"], np.array([2.5, 4.5]))

    def test_to_numpy_bool(self):
        data = [{"a": True, "b": False}, {"a": True, "b": False}]
        tbl = Table(data)
        view = tbl.view()
        v = view.to_numpy()
        assert np.array_equal(v["a"], np.array([True, True]))
        assert np.array_equal(v["b"], np.array([False, False]))

    def test_to_numpy_date(self):
        today = date.today()
        dt = datetime(today.year, today.month, today.day)
        data = [{"a": dt, "b": 2}, {"a": dt, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        v = view.to_numpy()
        assert np.array_equal(v["a"], np.array([dt, dt]))

    def test_to_numpy_datetime(self):
        dt = datetime(2019, 3, 15, 20, 30, 59, 6000)
        data = [{"a": dt}, {"a": dt}]
        tbl = Table(data)
        view = tbl.view()
        v = view.to_numpy()
        assert np.array_equal(v["a"], np.array([dt, dt]))

    def test_to_numpy_string(self):
        data = [{"a": "string1", "b": "string2"}, {"a": "string3", "b": "string4"}]
        tbl = Table(data)
        view = tbl.view()
        v = view.to_numpy()
        assert np.array_equal(v["a"], np.array(["string1", "string3"]))
        assert np.array_equal(v["b"], np.array(["string2", "string4"]))

    def test_to_numpy_none(self):
        data = [{"a": None, "b": None}, {"a": None, "b": None}]
        tbl = Table(data)
        view = tbl.view()
        v = view.to_numpy()
        assert np.array_equal(v["a"], np.array([None, None]))
        assert np.array_equal(v["b"], np.array([None, None]))

    def test_to_numpy_one(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"]
        )
        v = view.to_numpy()
        assert np.array_equal(v["__ROW_PATH__"], [[], [1]])
        assert np.array_equal(v["a"], np.array([2, 2]))
        assert np.array_equal(v["b"], np.array([4, 4]))

    def test_to_numpy_two(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"]
        )
        v = view.to_numpy()
        assert np.array_equal(v["__ROW_PATH__"], [[], [1]])
        assert np.array_equal(v["2|a"], np.array([2, 2]))
        assert np.array_equal(v["2|b"], np.array([4, 4]))

    def test_to_numpy_column_only(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            column_pivots=["b"]
        )
        v = view.to_numpy()
        assert np.array_equal(v["2|a"], np.array([1, 1]))
        assert np.array_equal(v["2|b"], np.array([2, 2]))

    def test_to_pandas_df_simple(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        df = pd.DataFrame(data)
        tbl = Table(df)
        view = tbl.view()
        df2 = view.to_df()
        if six.PY2:
            # order not guaranteed if columns are inferred using aggregate dict keys
            for col in df2.columns:
                assert col in ["index", "a", "b"]
        else:
            assert np.array_equal(df2.columns, pd.Index(["index", "a", "b"], dtype=object))

        assert np.array_equal(df2["a"].values, df["a"].values)
        assert np.array_equal(df2["b"].values, df["b"].values)

    def test_to_pandas_df_simple_series(self):
        inp = pd.Series([1, 2, 3], name="a")
        df = pd.DataFrame()
        df["a"] = pd.Series([1, 2, 3])
        tbl = Table(inp)
        view = tbl.view()
        df2 = view.to_df()
        if six.PY2:
            # order not guaranteed if columns are inferred using aggregate dict keys
            for col in df2.columns:
                assert col in ["index", "a"]
        else:
            assert np.array_equal(df2.columns, pd.Index(["index", "a"], dtype=object))
        assert np.array_equal(df2["a"].values, df["a"].values)

    # start_row/end_row
    def test_to_records_zero_over_max_row(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            end_row=1000
        )
        assert records == data

    def test_to_records_one_over_max_row(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"]
        )
        records = view.to_records(
            end_row=1000
        )
        assert records == [
            {'__ROW_PATH__': [], 'a': 5, 'b': 7},
            {'__ROW_PATH__': [1.5], 'a': 1.5, 'b': 2.5},
            {'__ROW_PATH__': [3.5], 'a': 3.5, 'b': 4.5}
        ]

    def test_to_records_two_over_max_row(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"]
        )
        records = view.to_records(
            end_row=1000
        )
        assert records == [
            {'2|a': 1, '2|b': 2, '4|a': 3, '4|b': 4, '__ROW_PATH__': []},
            {'2|a': 1, '2|b': 2, '4|a': None, '4|b': None, '__ROW_PATH__': [1]},
            {'2|a': None, '2|b': None, '4|a': 3, '4|b': 4, '__ROW_PATH__': [3]}
        ]

    def test_to_records_start_row(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_row=1
        )
        assert records == [{"a": 3, "b": 4}]

    def test_to_records_end_row(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            end_row=1
        )
        assert records == [{"a": 1, "b": 2}]

    def test_to_records_start_row_end_row(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_row=1,
            end_row=2
        )
        assert records == [{"a": 3, "b": 4}]

    def test_to_records_start_row_end_row_equiv(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_row=1,
            end_row=1
        )
        assert records == []

    def test_to_records_floor_start_row(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_row=1.5
        )
        assert records == [{"a": 3, "b": 4}]

    def test_to_records_ceil_end_row(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            end_row=0.5
        )
        assert records == [{"a": 1, "b": 2}]

    def test_to_records_floor_start_row_ceil_end_row(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_row=1.5,
            end_row=1.5
        )
        assert records == [{"a": 3, "b": 4}]

    def test_to_records_floor_start_row_ceil_end_row_equiv(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_row=1.5,
            end_row=0.5
        )
        assert records == []

    # start_col/end_col

    def test_to_records_zero_over_max_col(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            end_col=1000
        )
        assert records == data

    def test_to_records_zero_start_gt_end_col(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_col=2,
            end_col=1
        )
        assert records == [{}, {}]

    def test_to_records_zero_start_eq_end_col(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_col=1,
            end_col=1
        )
        assert records == [{}, {}]

    def test_to_records_one_over_max_col(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"]
        )
        records = view.to_records(
            end_col=1000
        )
        assert records == [
            {'__ROW_PATH__': [], 'a': 5, 'b': 7},
            {'__ROW_PATH__': [1.5], 'a': 1.5, 'b': 2.5},
            {'__ROW_PATH__': [3.5], 'a': 3.5, 'b': 4.5}
        ]

    def test_to_records_one_start_gt_end_col(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"]
        )
        records = view.to_records(
            start_col=2,
            end_col=1
        )
        assert records == [{}, {}, {}]

    def test_to_records_one_start_gt_end_col_large(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"]
        )
        records = view.to_records(
            start_col=20,
            end_col=19
        )
        assert records == [{}, {}, {}]

    def test_to_records_one_start_eq_end_col(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"]
        )
        records = view.to_records(
            start_col=0,
            end_col=0
        )
        assert records == [{'__ROW_PATH__': []}, {'__ROW_PATH__': [1.5]}, {'__ROW_PATH__': [3.5]}]

    def test_to_records_two_over_max_col(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"]
        )
        records = view.to_records(
            end_col=1000
        )
        assert records == [
            {'2|a': 1, '2|b': 2, '4|a': 3, '4|b': 4, '__ROW_PATH__': []},
            {'2|a': 1, '2|b': 2, '4|a': None, '4|b': None, '__ROW_PATH__': [1]},
            {'2|a': None, '2|b': None, '4|a': 3, '4|b': 4, '__ROW_PATH__': [3]}
        ]

    def test_to_records_start_col(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_col=1
        )
        assert records == [{"b": 2}, {"b": 4}]

    def test_to_records_end_col(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            end_col=1
        )
        assert records == [{"a": 1}, {"a": 3}]

    def test_to_records_two_end_col(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"]
        )
        records = view.to_records(
            end_row=12,
            end_col=5
        )
        assert records == [
            {'2|a': 1, '2|b': 2, '4|a': 3, '4|b': 4, '__ROW_PATH__': []},
            {'2|a': 1, '2|b': 2, '4|a': None, '4|b': None, '__ROW_PATH__': [1]},
            {'2|a': None, '2|b': None, '4|a': 3, '4|b': 4, '__ROW_PATH__': [3]}
        ]

    def test_to_records_two_start_gt_end_col(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"]
        )
        records = view.to_records(
            end_row=12,
            start_col=5,
            end_col=4
        )
        assert records == [{}, {}, {}]

    def test_to_records_two_start_gt_end_col_large_overage(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"]
        )
        records = view.to_records(
            end_row=12,
            start_col=50,
            end_col=49
        )
        assert records == [{}, {}, {}]

    def test_to_records_two_start_end_col_equiv(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"]
        )
        records = view.to_records(
            end_row=12,
            start_col=5,
            end_col=5
        )
        assert records == [{}, {}, {}]

    def test_to_records_two_sorted_start_gt_end_col(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"],
            sort=[["a", "desc"]]
        )
        records = view.to_records(
            end_row=12,
            start_col=5,
            end_col=4
        )
        assert records == [{}, {}, {}]

    def test_to_records_two_sorted_start_gt_end_col_large_overage(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"],
            sort=[["a", "desc"]]
        )
        records = view.to_records(
            end_row=12,
            start_col=20,
            end_col=30
        )
        assert records == [{}, {}, {}]

    def test_to_records_two_sorted_start_gt_end_col_overage(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            columns=[],
            row_pivots=["a"],
            column_pivots=["b"],
            sort=[["a", "desc"]]
        )
        records = view.to_records(
            end_row=12,
            start_col=1,
            end_col=3
        )
        assert records == [{}, {}, {}]

    def test_to_records_two_sorted_start_end_col(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"],
            sort=[["a", "desc"]]
        )
        records = view.to_records(
            start_col=1,
            end_col=2
        )
        assert records == [
            {'2|b': 2, '__ROW_PATH__': []},
            {'2|b': None, '__ROW_PATH__': [3]},
            {'2|b': 2, '__ROW_PATH__': [1]}
        ]

    def test_to_records_two_sorted_start_end_col_equiv(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"],
            sort=[["a", "desc"]]
        )
        records = view.to_records(
            end_row=12,
            start_col=5,
            end_col=5
        )
        assert records == [{}, {}, {}]

    def test_to_records_start_col_end_col(self):
        data = [{"a": 1, "b": 2, "c": 3}, {"a": 3, "b": 4, "c": 5}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_col=1,
            end_col=2
        )
        # start_col and end_col access columns at that index - dict key order not guaranteed in python2
        if six.PY2:
            # in this test, column c comes before b
            assert records == [{"c": 3}, {"c": 5}]
            # assert that in the general dataset, just to see if it holds true
            assert view.to_records() == [{"a": 1, "c": 3, "b": 2}, {"a": 3, "c": 5, "b": 4}]
        else:
            assert records == [{"b": 2}, {"b": 4}]

    def test_to_records_start_col_end_col_equiv(self):
        data = [{"a": 1, "b": 2, "c": 3}, {"a": 3, "b": 4, "c": 5}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_col=1,
            end_col=1
        )
        assert records == [{}, {}]

    def test_to_records_floor_start_col(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_col=1.5
        )
        assert records == [{"b": 2}, {"b": 4}]

    def test_to_records_ceil_end_col(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            end_col=1
        )
        assert records == [{"a": 1}, {"a": 3}]

    def test_to_records_two_ceil_end_col(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"]
        )
        records = view.to_records(
            end_row=12,
            end_col=4.5
        )
        assert records == [
            {'2|a': 1, '2|b': 2, '4|a': 3, '4|b': 4, '__ROW_PATH__': []},
            {'2|a': 1, '2|b': 2, '4|a': None, '4|b': None, '__ROW_PATH__': [1]},
            {'2|a': None, '2|b': None, '4|a': 3, '4|b': 4, '__ROW_PATH__': [3]}
        ]

    def test_to_records_floor_start_col_ceil_end_col(self):
        data = [{"a": 1, "b": 2, "c": 3}, {"a": 3, "b": 4, "c": 5}]
        tbl = Table(data)
        view = tbl.view()
        records = view.to_records(
            start_col=1.5,
            end_col=1.5
        )
        # start_col and end_col access columns at that index - dict key order not guaranteed in python2
        if six.PY2:
            # in this test, column c comes before b
            assert records == [{"c": 3}, {"c": 5}]
            # assert that in the general dataset, just to see if it holds true
            assert view.to_records() == [{"a": 1, "c": 3, "b": 2}, {"a": 3, "c": 5, "b": 4}]
        else:
            assert records == [{"b": 2}, {"b": 4}]

    def test_to_dict_start_col_end_col(self):
        data = [{"a": 1, "b": 2, "c": 3, "d": 4}, {"a": 3, "b": 4, "c": 5, "d": 6}]
        tbl = Table(data)
        view = tbl.view()
        d = view.to_dict(
            start_col=1,
            end_col=3
        )
        assert d == {
            "b": [2, 4],
            "c": [3, 5]
        }

    # to csv

    def test_to_csv_symmetric(self):
        csv = "a,b\n1,2\n3,4"
        if six.PY2:
            csv = unicode(csv)
        df = pd.read_csv(StringIO(csv))
        tbl = Table(df)
        view = tbl.view()
        if six.PY2:
            if IS_WIN:
                assert view.to_csv() == ",a,b,index\r\n0,1,2,0\r\n1,3,4,1\r\n"
            else:
                assert view.to_csv() == ",a,b,index\n0,1,2,0\n1,3,4,1\n"
        else:
            if IS_WIN:
                assert view.to_csv() == ",index,a,b\r\n0,0,1,2\r\n1,1,3,4\r\n"
            else:
                assert view.to_csv() == ",index,a,b\n0,0,1,2\n1,1,3,4\n"

    def test_to_csv_int(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv() == ",a,b\r\n0,1,2\r\n1,3,4\r\n"
        else:
            assert view.to_csv() == ",a,b\n0,1,2\n1,3,4\n"

    def test_to_csv_float(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv() == ",a,b\r\n0,1.5,2.5\r\n1,3.5,4.5\r\n"
        else:
            assert view.to_csv() == ",a,b\n0,1.5,2.5\n1,3.5,4.5\n"

    def test_to_csv_date(self):
        today = date.today()
        dt_str = today.strftime("%Y-%m-%d")
        data = [{"a": today, "b": 2}, {"a": today, "b": 4}]
        tbl = Table(data)
        assert tbl.schema()["a"] == date
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv() == ",a,b\r\n0,{},2\r\n1,{},4\r\n".format(dt_str, dt_str)
        else:
            assert view.to_csv() == ",a,b\n0,{},2\n1,{},4\n".format(dt_str, dt_str)

    @mark.skip(reason="pandas does not support date-specific formated in the date_format kwarg")
    def test_to_csv_date_formatted(self):
        today = date.today()
        dt_str = today.strftime("%Y/%m/%d")
        data = [{"a": today, "b": 2}, {"a": today, "b": 4}]
        tbl = Table(data)
        assert tbl.schema()["a"] == date
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv(formatted=True) == ",a,b\r\n0,{},2\r\n1,{},4\r\n".format(dt_str, dt_str)
        else:
            assert view.to_csv(formatted=True) == ",a,b\n0,{},2\n1,{},4\n".format(dt_str, dt_str)

    def test_to_csv_datetime(self):
        dt = datetime(2019, 3, 15, 20, 30, 59, 6000)
        dt_str = dt.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
        data = [{"a": dt, "b": 2}, {"a": dt, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv() == ",a,b\r\n0,{},2\r\n1,{},4\r\n".format(dt_str, dt_str)
        else:
            assert view.to_csv() == ",a,b\n0,{},2\n1,{},4\n".format(dt_str, dt_str)

    def test_to_csv_datetime_formatted(self):
        dt = datetime(2019, 3, 15, 20, 30, 59, 6000)
        dt_str = dt.strftime("%Y/%m/%d %H:%M:%S")
        data = [{"a": dt, "b": 2}, {"a": dt, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv(formatted=True) == ",a,b\r\n0,{},2\r\n1,{},4\r\n".format(dt_str, dt_str)
        else:
            assert view.to_csv(formatted=True) == ",a,b\n0,{},2\n1,{},4\n".format(dt_str, dt_str)

    def test_to_csv_bool(self):
        data = [{"a": True, "b": False}, {"a": True, "b": False}]
        tbl = Table(data)
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv() == ",a,b\r\n0,True,False\r\n1,True,False\r\n"
        else:
            assert view.to_csv() == ",a,b\n0,True,False\n1,True,False\n"

    def test_to_csv_string(self):
        data = [{"a": "string1", "b": "string2"}, {"a": "string3", "b": "string4"}]
        tbl = Table(data)
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv() == ",a,b\r\n0,string1,string2\r\n1,string3,string4\r\n"
        else:
            assert view.to_csv() == ",a,b\n0,string1,string2\n1,string3,string4\n"

    def test_to_csv_none(self):
        data = [{"a": None, "b": None}, {"a": None, "b": None}]
        tbl = Table(data)
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv() == ",a,b\r\n0,,\r\n1,,\r\n"
        else:
            assert view.to_csv() == ",a,b\n0,,\n1,,\n"

    def test_to_csv_custom_rows(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv(start_row=1) == ",a,b\r\n0,3,4\r\n"
        else:
            assert view.to_csv(start_row=1) == ",a,b\n0,3,4\n"

    def test_to_csv_custom_cols(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv(start_col=1) == ",b\r\n0,2\r\n1,4\r\n"
        else:
            assert view.to_csv(start_col=1) == ",b\n0,2\n1,4\n"

    def test_to_csv_custom_rows_cols(self):
        data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}]
        tbl = Table(data)
        view = tbl.view()
        if IS_WIN:
            assert view.to_csv(start_row=1, start_col=1) == ",b\r\n0,4\r\n"
        else:
            assert view.to_csv(start_row=1, start_col=1) == ",b\n0,4\n"

    def test_to_csv_one(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"]
        )
        if six.PY2:
            if IS_WIN:
                assert view.to_csv() == ",__ROW_PATH__,a,b\r\n0,[],2,4\r\n1,[1],2,4\r\n"
            else:
                assert view.to_csv() == ",__ROW_PATH__,a,b\n0,[],2,4\n1,[1],2,4\n"
        else:
            if IS_WIN:
                assert view.to_csv() == ",__ROW_PATH__,a,b\r\n0,[],2,4\r\n1,[1],2,4\r\n"
            else:
                assert view.to_csv() == ",__ROW_PATH__,a,b\n0,[],2,4\n1,[1],2,4\n"

    def test_to_csv_two(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"]
        )
        if six.PY2:
            if IS_WIN:
                assert view.to_csv() == ",2|a,2|b,__ROW_PATH__\r\n0,2,4,[]\r\n1,2,4,[1]\r\n"
            else:
                assert view.to_csv() == ",2|a,2|b,__ROW_PATH__\n0,2,4,[]\n1,2,4,[1]\n"
        else:
            if IS_WIN:
                assert view.to_csv() == ",__ROW_PATH__,2|a,2|b\r\n0,[],2,4\r\n1,[1],2,4\r\n"
            else:
                assert view.to_csv() == ",__ROW_PATH__,2|a,2|b\n0,[],2,4\n1,[1],2,4\n"

    def test_to_csv_column_only(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            column_pivots=["b"]
        )
        if IS_WIN:
            assert view.to_csv() == ",2|a,2|b\r\n0,1,2\r\n1,1,2\r\n"
        else:
            assert view.to_csv() == ",2|a,2|b\n0,1,2\n1,1,2\n"

    def test_to_csv_one_no_columns(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            columns=[]
        )
        if six.PY2:
            if IS_WIN:
                assert view.to_csv() == ",__ROW_PATH__\r\n0,[]\r\n1,[1]\r\n"
            else:
                assert view.to_csv() == ",__ROW_PATH__\n0,[]\n1,[1]\n"
        else:
            if IS_WIN:
                assert view.to_csv() == ",__ROW_PATH__\r\n0,[]\r\n1,[1]\r\n"
            else:
                assert view.to_csv() == ",__ROW_PATH__\n0,[]\n1,[1]\n"

    def test_to_csv_two_no_columns(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            row_pivots=["a"],
            column_pivots=["b"],
            columns=[]
        )
        if six.PY2:
            if IS_WIN:
                assert view.to_csv() == ",__ROW_PATH__\r\n0,[]\r\n1,[1]\r\n"
            else:
                assert view.to_csv() == ",__ROW_PATH__\n0,[]\n1,[1]\n"
        else:
            if IS_WIN:
                assert view.to_csv() == ",__ROW_PATH__\r\n0,[]\r\n1,[1]\r\n"
            else:
                assert view.to_csv() == ",__ROW_PATH__\n0,[]\n1,[1]\n"

    def test_to_csv_column_only_no_columns(self):
        data = [{"a": 1, "b": 2}, {"a": 1, "b": 2}]
        tbl = Table(data)
        view = tbl.view(
            column_pivots=["b"],
            columns=[]
        )

        if IS_WIN:
            assert view.to_csv() == '""\r\n'
        else:
            assert view.to_csv() == '""\n'

    # implicit index

    def test_to_format_implicit_index_records(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records(index=True) == [
            {"__INDEX__": [0], "a": 1.5, "b": 2.5},
            {"__INDEX__": [1], "a": 3.5, "b": 4.5}
        ]

    def test_to_format_implicit_index_dict(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_dict(index=True) == {
            "__INDEX__": [[0], [1]],
            "a": [1.5, 3.5],
            "b": [2.5, 4.5]
        }

    def test_to_format_implicit_id_records(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_records(id=True) == [
            {"__ID__": [0], "a": 1.5, "b": 2.5},
            {"__ID__": [1], "a": 3.5, "b": 4.5}
        ]

    def test_to_format_implicit_id_dict(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        assert view.to_dict(id=True) == {
            "__ID__": [[0], [1]],
            "a": [1.5, 3.5],
            "b": [2.5, 4.5]
        }

    def test_to_format_implicit_index_two_dict(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view(row_pivots=["a"], column_pivots=["b"])
        assert view.to_dict(index=True) == {
            '2.5|a': [1.5, 1.5, None],
            '2.5|b': [2.5, 2.5, None],
            '4.5|a': [3.5, None, 3.5],
            '4.5|b': [4.5, None, 4.5],
            '__INDEX__': [[], [], []],  # index needs to be the same length as each column
            '__ROW_PATH__': [[], [1.5], [3.5]]
        }

    def test_to_format_implicit_index_two_dict(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view(row_pivots=["a"], column_pivots=["b"])
        assert view.to_dict(id=True) == {
            '2.5|a': [1.5, 1.5, None],
            '2.5|b': [2.5, 2.5, None],
            '4.5|a': [3.5, None, 3.5],
            '4.5|b': [4.5, None, 4.5],
            '__ID__': [[], [1.5], [3.5]],  # index needs to be the same length as each column
            '__ROW_PATH__': [[], [1.5], [3.5]]
        }

    def test_to_format_implicit_index_np(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data)
        view = tbl.view()
        cols = view.to_numpy(index=True)
        assert np.array_equal(cols["__INDEX__"], np.array([[0], [1]]))

    def test_to_format_explicit_index_records(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data, index="a")
        view = tbl.view()
        assert view.to_records(index=True) == [
            {"__INDEX__": [1.5], "a": 1.5, "b": 2.5},
            {"__INDEX__": [3.5], "a": 3.5, "b": 4.5}
        ]

    def test_to_format_explicit_index_dict(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data, index="a")
        view = tbl.view()
        assert view.to_dict(index=True) == {
            "__INDEX__": [[1.5], [3.5]],
            "a": [1.5, 3.5],
            "b": [2.5, 4.5]
        }

    def test_to_format_explicit_index_np(self):
        data = [{"a": 1.5, "b": 2.5}, {"a": 3.5, "b": 4.5}]
        tbl = Table(data, index="a")
        view = tbl.view()
        cols = view.to_numpy(index=True)
        assert np.array_equal(cols["__INDEX__"], np.array([[1.5], [3.5]]))

    def test_to_format_explicit_index_str_records(self):
        data = [{"a": "a", "b": 2.5}, {"a": "b", "b": 4.5}]
        tbl = Table(data, index="a")
        view = tbl.view()
        assert view.to_records(index=True) == [
            {"__INDEX__": ["a"], "a": "a", "b": 2.5},
            {"__INDEX__": ["b"], "a": "b", "b": 4.5}
        ]

    def test_to_format_explicit_index_datetime_records(self):
        data = [{"a": datetime(2019, 7, 11, 9, 0), "b": 2.5}, {"a": datetime(2019, 7, 11, 9, 1), "b": 4.5}]
        tbl = Table(data, index="a")
        view = tbl.view()
        assert view.to_records(index=True) == [
            {"__INDEX__": [datetime(2019, 7, 11, 9, 0)], "a": datetime(2019, 7, 11, 9, 0), "b": 2.5},
            {"__INDEX__": [datetime(2019, 7, 11, 9, 1)], "a": datetime(2019, 7, 11, 9, 1), "b": 4.5}
        ]

max_stars_repo_name: sunpar/perspective

max_issues_repo_name: sunpar/perspective

max_forks_repo_name: sunpar/perspective


text:
import pygame
import time
import itertools
import numpy as np
import numpy.linalg as la

from ball import Ball
from wall import Wall

class Game:
    def __init__(self, width, height):
        pygame.init()
        self.screen = pygame.display.set_mode((width, height))
        self.width = width
        self.height = height
        self.frames_per_second = 60
        self.max_num_balls = 30
        self.gravity = np.array([0, 300], dtype=np.float)

    def update(self, dt):
        """
        This function is called onces per frame.
        We want to compute forces and detect collisions for all of our balls.
        """
        # TODO: 11. Uncomment this when we tell you.
        # for ball_a, ball_b in itertools.combinations(self.balls, r=2):
        #     Ball.ball_ball_collision(ball_a, ball_b, dt)

        # TODO: 5. Use a for loop to apply the force of gravity to each ball.
        # HINT: force_gravity = acceleration_gravity * mass

        # TODO: 8. Use a for loop to add wall collisions.

        for ball in self.balls:
            ball.update(dt)

    def render(self):
        """
        Draws all of our objects on to the screen.
        """
        self.screen.fill((0, 0, 0))
        for _object in self.balls + self.walls:
            _object.draw(self.screen)
        pygame.display.flip()

    def play(self):
        """
        Sets up our scene and executes the game loop.
        """
        self.is_running = True
        self.is_building_wall = False
        self.balls = []
        color = (255, 255, 255)
        self.walls = [Wall((1, 1), (self.width-1, 1), color),                           # ceiling
                      Wall((1, 1), (1, self.height-1), color),                          # left wall
                      Wall((1, self.height-1), (self.width-1, self.height-1), color),   # floor
                      Wall((self.width-1, 1), (self.width-1, self.height-1), color)]    # right wall

        last_frame_time = time.time()
        while self.is_running:
            time_remaining = 1. / self.frames_per_second - (time.time() - last_frame_time)
            if time_remaining > 0:
                time.sleep(time_remaining)
            current_frame_time = time.time()
            dt = current_frame_time - last_frame_time
            dt = max(0.00001, min(dt, 0.1))
            last_frame_time = current_frame_time

            self.event_handler()
            self.update(dt)
            self.render()

    def event_handler(self):
        for event in pygame.event.get():
            if event.type is pygame.QUIT:
                self.is_running = False
            elif event.type is pygame.KEYDOWN:
                if event.key is pygame.K_ESCAPE:
                    self.is_running = False
                elif event.key is pygame.K_RETURN:
                    if len(self.balls) < self.max_num_balls:
                        self.add_ball()
            elif event.type is pygame.MOUSEBUTTONDOWN:
                if event.button is 1:   # Left Mouse Button
                    self.is_building_wall = True
                    pos = pygame.mouse.get_pos()
                    self.add_wall(pos, pos)
            elif event.type is pygame.MOUSEBUTTONUP:
                if event.button is 1:   # Left Mouse Button
                    self.is_building_wall = False
            elif event.type is pygame.MOUSEMOTION:
                if self.is_building_wall:
                    pos = pygame.mouse.get_pos()
                    self.walls[-1].b = pos

    def add_ball(self):
        """
        Add a new ball object to our scene.
        """
        radius = int(np.random.randint(10, 30))
        x = int(np.random.randint(radius, self.width - radius))
        y = radius + 5
        color = np.floor(np.random.random(3) * 200 + 55).astype(np.int)
        self.balls.append(Ball((x, y), radius, color))

    def add_wall(self, a, b):
        """
        Add a new wall object to our scene.
        """
        color = np.floor(np.random.rand(3) * 200 + 55).astype(np.int)
        self.walls.append(Wall(a, b, color))

if __name__ == "__main__":
    mygame = Game(width=400, height=400)
    mygame.play()

max_stars_repo_name: dnagarkar/bubble-deflector

max_issues_repo_name: dnagarkar/bubble-deflector

max_forks_repo_name: dnagarkar/bubble-deflector


text:
"""
===========================================
Plane wave reflection from periodic surface
===========================================
Common framework for scatter with a plane wave source and a periodic surface
"""

import numpy as np
from math import pi

class Bragg:
    """Compute reflection coefficents for cosine surface"""
    def __init__(self, Lper, c=1500., attn=0):
        """surface specific properties
        attn: volume attenuation, dB / km"""
        self.Lper = Lper
        self.Kper = 2 * pi / Lper
        self.c = c
        self.attn = attn
        self.delta = lambda k: 1j * self.attn / 8686.\
                               / np.real(k + np.spacing(1))

    def kacous(self, facous):
        """complex wavenumeber of medium"""
        k = 2 * pi * facous / self.c
        return k + 1j * self.delta(k)

    def xsampling(self, facous, decimation=8):
        """Make a periodically sampled xaxis"""
        dx = self.c / (decimation * facous)
        numx = int(np.ceil(self.L / dx))
        dx = self.L / numx
        xaxis = np.arange(numx) * dx
        return xaxis, dx

    def qvec(self, theta_inc, num_eva, facous):
        """Return vector of bragg grating orders
        cutoff after num_eva evanescent orders on each side
        """
        kacous = np.real(self.kacous(facous))
        kx = np.real(np.cos(theta_inc) * kacous)
        num_p = np.fix((kacous - kx) / self.Kper) + num_eva
        num_n = np.fix((kacous + kx) / self.Kper) + num_eva
        qvec = np.arange(-num_n, num_p + 1)
        return qvec

    def bragg_angles(self, theta_inc, qs, facous):
        """Computer the brag angle cosine vectors"""
        kacous = self.kacous(facous)

        # compute bragg orders
        a0 = np.real(np.cos(theta_inc) * kacous)
        b0 = np.conj(np.sqrt(kacous ** 2 - a0 ** 2))
        aq = a0 + qs * self.Kper
        bq = np.conj(np.sqrt(kacous ** 2 - aq ** 2))
        return a0, aq, b0, bq

    def p_sca(self, theta_inc, qs, facous, rs, xsrc, zsrc, xrcr, zrcr):
        """
        Scattered pressure field from plane wave reflection coefficents
        """
        a0, aq, b0, bq = self.bragg_angles(theta_inc, qs, facous)
        phase = -a0 * xsrc - b0 * zsrc + aq * xrcr - bq * zrcr
        p_sca = rs @ np.exp(-1j * phase)
        return p_sca

    def r_energy(self, theta_inc, qs, facous, rs):
        """Calculate the energy conservation relative to 1"""
        kacous = self.kacous(facous)
        _, aq, b0, bq = self.bragg_angles(theta_inc, qs, facous)
        # compute energy
        reali = np.abs(np.real(aq ** 2)) <= np.real(kacous) ** 2
        en_conn = np.abs(rs[reali]) ** 2 * np.real(bq[reali]) / np.real(b0)
        return np.sum(en_conn)