text:
# -*- coding: utf-8 -*-
"""
Created on Tue Nov 21 17:35:53 2017

@author: bokorn
"""

import numpy as np
import cv2

from se3_distributions.utils.image_preprocessing import unprocessImages
from se3_distributions.eval.hyper_distance import ExemplarDistPoseEstimator
from quat_math import (quatAngularDiff,
                       quat2AxisAngle,
                       quaternion_about_axis,
                       quaternion_multiply,
                       quaternion_inverse,
                       quaternion_from_matrix)

from se3_distributions.utils import to_np

delta_quat = np.array([.5,.5,.5,.5])
#flip_x = np.array([-1,1,1,1])
#def convertQuat(q):
#    q_flip = quaternion_inverse(q.copy())
#    q_flip[2] *= -1
#    return quaternion_multiply(delta_quat, q_flip)

def evaluateDistanceNetwork(estimator, data_loader, trans_mat=np.eye(4),
                            inverse_distance = True,
                            use_converter = False):
    trans_quat = quaternion_inverse(quaternion_from_matrix(trans_mat))
    if(use_converter):
        def convertQuat(q):
            q_flip = quaternion_inverse(q.copy())
            q_flip = quaternion_multiply(trans_quat, q_flip)
            q_flip[2] *= -1
            return quaternion_multiply(delta_quat, q_flip)
    else:
        def convertQuat(q):
            return q

    def trueDiff(q):
        q_adj = convertQuat(q)
        true_diff = []
        for v in estimator.base_vertices:
            true_diff.append(quatAngularDiff(q_adj, v))
        return np.array(true_diff)

    if(inverse_distance):
        sign = -1.0
    else:
        sign = 1
    gt_ranks = []
    gt_outputs = []
    top_ranks = []
    top_dists = []
    angle_vec = []
    output_vec = []
    gt_pos_images = []
    gt_neg_images = []
    top_pos_images = []
    top_neg_images = []

    for j, (imgs, _, quats, _, _) in enumerate(data_loader):
        diff = estimator.estimate(imgs[0], preprocess=False)
        diff = to_np(diff.detach())
        true_dist = trueDiff(to_np(quats[0][0]))
        top_idx = np.argmin(sign*diff)
        true_idx = np.argmin(true_dist)

        angle_vec.extend(true_dist.tolist())
        output_vec.extend(diff.tolist())

        rank_gt = np.nonzero(np.argsort(sign*diff) == true_idx)[0][0]
        rank_top = np.nonzero(np.argsort(true_dist) == top_idx)[0][0]
        output_gt = diff[true_idx]
        dist_top = true_dist[top_idx]*180/np.pi
        gt_ranks.append(rank_gt)
        gt_outputs.append(output_gt)
        top_ranks.append(rank_top)
        top_dists.append(dist_top)
        #print('True ranking: {}'.format(np.nonzero(np.argsort(-diff) == true_idx)[0][0]))
        #print('Top scored ranking: {}'.format(np.nonzero(np.argsort(true_dist) == top_idx)[0][0]))
        #print('Top distance: {}'.format(true_dist[top_idx]*180/np.pi))
        if(rank_gt < 5):
            top_img = unprocessImages(estimator.base_renders[top_idx:top_idx+1])[0]
            tgt_img = unprocessImages(imgs[0])[0]
            gt_img  = unprocessImages(estimator.base_renders[true_idx:true_idx+1])[0]
            gt_pos_images = [top_img, tgt_img, gt_img, j]
        elif(rank_gt > 55):
            top_img = unprocessImages(estimator.base_renders[top_idx:top_idx+1])[0]
            tgt_img = unprocessImages(imgs[0])[0]
            gt_img  = unprocessImages(estimator.base_renders[true_idx:true_idx+1])[0]
            gt_neg_images = [top_img, tgt_img, gt_img, j]
        elif(rank_top < 5):
            top_img = unprocessImages(estimator.base_renders[top_idx:top_idx+1])[0]
            tgt_img = unprocessImages(imgs[0])[0]
            gt_img  = unprocessImages(estimator.base_renders[true_idx:true_idx+1])[0]
            top_pos_images = [top_img, tgt_img, gt_img, j]
        elif(rank_top > 5):
            top_img = unprocessImages(estimator.base_renders[top_idx:top_idx+1])[0]
            tgt_img = unprocessImages(imgs[0])[0]
            gt_img  = unprocessImages(estimator.base_renders[true_idx:true_idx+1])[0]
            top_neg_images = [top_img, tgt_img, gt_img, j]

        #input("Press Enter to continue...")
        #import IPython; IPython.embed()
    images = {'gt_pos':gt_pos_images,
              'gt_neg':gt_neg_images,
              'top_pos':top_pos_images,
              'top_neg':top_neg_images}

    data = np.array([angle_vec, output_vec])
    return np.array(gt_ranks), np.array(top_ranks),\
           np.array(gt_outputs), np.array(top_dists),\
           data, images

def getArgs():
    from argparse import ArgumentParser

    parser = ArgumentParser()
    parser.add_argument('--results_prefix', type=str)
    parser.add_argument('--weight_file', type=str)
    parser.add_argument('--data_folder', type=str)
    parser.add_argument('--dataset_type', type=str, default='numpy')
    parser.add_argument('--background_data_file', type=str, default=None)

    parser.add_argument('--model_filename', type=str, default=None)

    parser.add_argument('--model_type', type=str, default='alexnet')
    parser.add_argument('--compare_type', type=str, default='sigmoid')
    parser.add_argument('--loss_type', type=str, default='exp')

    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--num_workers', type=int, default=4)

    parser.add_argument('--num_samples', type=int, default=None)
    parser.add_argument('--inverse_distance', dest='inverse_distance', action='store_true')

    parser.add_argument('--base_level', type=int, default=0)

    args = parser.parse_args()
    if(args.background_data_file is not None):
        with open(args.background_data_file, 'r') as f:
            background_filenames = f.read().split()
    else:
        background_filenames = None

    if(args.model_filename is None):
        if(args.data_folder[-1] == '/'):
            args.data_folder = args.data_folder[:-1]
        model_class, model_name = args.data_folder.split('/')[-2:]
        if(model_class == 'linemod'):
            args.model_filename = '/scratch/bokorn/data/benchmarks/linemod/' + \
                                  '{}/mesh.ply'.format(model_name)
        else:
            args.model_filename = '/scratch/bokorn/data/models/shapenetcore/' + \
                                  '{}/{}/model.obj'.format(model_class, model_name)

    return args

def main(weight_file,
         model_filename,
         data_folder,
         dataset_type,
         results_prefix,
         model_type = 'alexnet',
         compare_type = 'sigmoid',
         inverse_distance = True,
         background_filenames = None,
         base_level = 0,
         ):
    import os
    import time
    import cv2
    import torch
    from torch.utils.data import DataLoader
    from se3_distributions.models.pose_networks import gen_pose_net

    print(results_prefix)
    t = time.time()
    dist_net = gen_pose_net(model_type.lower(),
                            compare_type.lower(),
                            output_dim = 1,
                            pretrained = False)

    dist_net.load_state_dict(torch.load(weight_file))
    #print('Weights load time: {}s'.format(round(time.time()-t, 2)))

    if(dataset_type.lower() == 'numpy'):
        from se3_distributions.datasets.numpy_dataset import NumpyImageDataset as Dataset
    elif(dataset_type.lower() == 'linemod'):
        from se3_distributions.datasets.benchmark_dataset import LinemodDataset as Dataset
    elif(dataset_type.lower() == 'linemod_masked'):
        from functools import partial
        from se3_distributions.datasets.benchmark_dataset import LinemodDataset
        Dataset = partial(LinemodDataset, use_mask = True)
    else:
        raise ValueError('Dataset type {} not implemented'.format(data_type))

    t = time.time()
    data_loader = DataLoader(Dataset(data_folders=data_folder,
                                     img_size = (224, 224),
                                     max_orientation_offset = None,
                                     max_orientation_iters = None,
                                     model_filenames=None,
                                     background_filenames = background_filenames),
                             num_workers=0,
                             batch_size=1,
                             shuffle=True)

    data_loader.dataset.loop_truth = [1]
    #print('Dataset initialization time: {}s'.format(round(time.time()-t, 2)))

    if('cam/mesh.ply' in model_filename or
       'eggbox/mesh.ply' in model_filename):
        trans_filename = model_filename.replace('mesh.ply', 'transform.dat')
        with open(trans_filename, 'r') as f:
            trans_data = f.read()
        trans = np.array([float(line.split()[1]) for line in trans_data.splitlines()[1:]])
        trans_mat = np.vstack([trans.reshape((3,4)), np.array([0,0,0,1])])
    else:
        trans_mat = np.eye(4)

    estimator = ExemplarDistPoseEstimator(model_filename, dist_net, use_bpy_renderer = True, base_level=base_level)
    res = evaluateDistanceNetwork(estimator, data_loader, trans_mat,
                                  inverse_distance = inverse_distance,
                                  use_converter = 'linemod' in dataset_type.lower())
    results_dir = os.path.dirname(results_prefix)
    os.makedirs(results_dir, exist_ok=True)
    np.savez(results_prefix + 'distance.npz',
            gt_ranks = res[0],
            top_ranks = res[1],
            gt_outputs = res[2],
            top_dists = res[3],
            data = res[4])

    image_path = os.path.join(results_dir, 'images')
    os.makedirs(image_path, exist_ok=True)
    images = res[5]

    if(len(images['gt_pos']) > 0):
        data_idx = images['gt_pos'][3]
        cv2.imwrite(image_path + '/gt_pos_top_img_{}.png'.format(data_idx),
                    images['gt_pos'][0])
        cv2.imwrite(image_path + '/gt_pos_tgt_img_{}.png'.format(data_idx),
                    images['gt_pos'][1])
        cv2.imwrite(image_path + '/gt_pos_gt_img_{}.png'.format(data_idx),
                    images['gt_pos'][2])
        with open(image_path + '/gt_pos.txt', 'w') as f:
            f.write('gt_rank:   {}\n'.format(res[0][data_idx]))
            f.write('top_rank:  {}\n'.format(res[1][data_idx]))
            f.write('gt_output: {}\n'.format(res[2][data_idx]))
            f.write('top_dist:  {}\n'.format(res[3][data_idx]*180/np.pi))

    if(len(images['gt_neg']) > 0):
        data_idx = images['gt_neg'][3]
        cv2.imwrite(image_path + '/gt_neg_top_img_{}.png'.format(data_idx),
                    images['gt_neg'][0])
        cv2.imwrite(image_path + '/gt_neg_tgt_img_{}.png'.format(data_idx),
                    images['gt_neg'][1])
        cv2.imwrite(image_path + '/gt_neg_gt_img_{}.png'.format(data_idx),
                    images['gt_neg'][2])
        with open(image_path + '/gt_neg.txt', 'w') as f:
            f.write('gt_rank:   {}\n'.format(res[0][data_idx]))
            f.write('top_rank:  {}\n'.format(res[1][data_idx]))
            f.write('gt_output: {}\n'.format(res[2][data_idx]))
            f.write('top_dist:  {}\n'.format(res[3][data_idx]*180/np.pi))

    if(len(images['top_pos']) > 0):
        data_idx = images['top_pos'][3]
        cv2.imwrite(image_path + '/top_pos_top_img_{}.png'.format(data_idx),
                    images['top_pos'][0])
        cv2.imwrite(image_path + '/top_pos_tgt_img_{}.png'.format(data_idx),
                    images['top_pos'][1])
        cv2.imwrite(image_path + '/top_pos_gt_img_{}.png'.format(data_idx),
                    images['top_pos'][2])
        with open(image_path + '/top_pos.txt', 'w') as f:
            f.write('gt_rank:   {}\n'.format(res[0][data_idx]))
            f.write('top_rank:  {}\n'.format(res[1][data_idx]))
            f.write('gt_output: {}\n'.format(res[2][data_idx]))
            f.write('top_dist:  {}\n'.format(res[3][data_idx]*180/np.pi))

    if(len(images['top_neg']) > 0):
        data_idx = images['top_neg'][3]
        cv2.imwrite(image_path + '/top_neg_top_img_{}.png'.format(data_idx),
                    images['top_neg'][0])
        cv2.imwrite(image_path + '/top_neg_tgt_img_{}.png'.format(data_idx),
                    images['top_neg'][1])
        cv2.imwrite(image_path + '/top_neg_gt_img_{}.png'.format(data_idx),
                    images['top_neg'][2])
        with open(image_path + '/top_neg.txt', 'w') as f:
            f.write('gt_rank:   {}\n'.format(res[0][data_idx]))
            f.write('top_rank:  {}\n'.format(res[1][data_idx]))
            f.write('gt_output: {}\n'.format(res[2][data_idx]))
            f.write('top_dist:  {}\n'.format(res[3][data_idx]*180/np.pi))

    print('Mean GT Rank: {}'.format(res[0].mean()))
    print('Mean Top Rank: {}'.format(res[1].mean()))
    print('Mean Top Dist: {}'.format(res[2].mean()))

if __name__=='__main__':
    args = getArgs()
    main(weight_file = args.weight_file,
         model_filename = args.model_filename,
         data_folder = args.data_folder,
         dataset_type = args.dataset_type,
         results_prefix = args.results_prefix,
         model_type = args.model_type,
         compare_type = args.compare_type,
         background_filenames = args.background_filenames,
         base_level = args.base_level)


text:
import numpy as np

import eqdes.nonlinear_foundation
from tests import models_for_testing as ml
from eqdes import dbd
from eqdes import models as dm
from eqdes import design_spectra
import sfsimodels as sm
import geofound as gf

from tests.checking_tools import isclose

def test_ddbd_frame_fixed_small():

    hz = dm.Hazard()
    ml.load_hazard_test_data(hz)
    fb = ml.initialise_frame_building_test_data()
    frame_ddbd = dbd.design_rc_frame(fb, hz)

    assert isclose(frame_ddbd.delta_d, 0.2400, rel_tol=0.001), frame_ddbd.delta_d
    assert isclose(frame_ddbd.mass_eff, 67841.581, rel_tol=0.001), frame_ddbd.mass_eff
    assert isclose(frame_ddbd.height_eff, 14.34915, rel_tol=0.001), frame_ddbd.height_eff
    assert isclose(frame_ddbd.mu, 1.689, rel_tol=0.001), frame_ddbd.mu
    assert isclose(frame_ddbd.theta_y, 0.0099, rel_tol=0.001), frame_ddbd.theta_y
    assert isclose(frame_ddbd.xi, 0.123399, rel_tol=0.001), frame_ddbd.xi
    assert isclose(frame_ddbd.eta, 0.69867, rel_tol=0.001), frame_ddbd.eta
    assert isclose(frame_ddbd.t_eff, 2.09646, rel_tol=0.001), frame_ddbd.t_eff

def test_ddbd_frame_consistent():
    """
    Test the DBD of a fixed base frame is the same as the SFSI frame when the soil is very stiff.
    :return:
    """

    fb = ml.initialise_frame_building_test_data()
    hz = dm.Hazard()
    sl = dm.Soil()
    fd = dm.RaftFoundation()
    ml.load_hazard_test_data(hz)
    ml.load_soil_test_data(sl)
    ml.load_raft_foundation_test_data(fd)
    frame_ddbd = dbd.design_rc_frame(fb, hz)
    sl.override("g_mod", 1.0e10)  # make soil very stiff
    fd.height = 2.0  # add some height to the foundation
    frame_sfsi_dbd = dbd.design_rc_frame_w_sfsi_via_millen_et_al_2018(fb, hz, sl, fd, found_rot=1e-6)
    assert isclose(frame_sfsi_dbd.theta_f, 0.0, abs_tol=1e-5)

    assert isclose(frame_ddbd.delta_d, frame_sfsi_dbd.delta_d, rel_tol=0.01), frame_sfsi_dbd.delta_d
    assert isclose(frame_ddbd.mass_eff, frame_sfsi_dbd.mass_eff, rel_tol=0.001), frame_sfsi_dbd.mass_eff
    assert isclose(frame_ddbd.height_eff, frame_sfsi_dbd.height_eff - fd.height, rel_tol=0.001), frame_sfsi_dbd.height_eff
    assert isclose(frame_ddbd.mu, frame_sfsi_dbd.mu, rel_tol=0.001), frame_sfsi_dbd.mu
    assert isclose(frame_ddbd.theta_y, frame_sfsi_dbd.theta_y, rel_tol=0.001), frame_sfsi_dbd.theta_y
    assert isclose(frame_ddbd.xi, frame_sfsi_dbd.xi, rel_tol=0.001), frame_sfsi_dbd.xi
    assert isclose(frame_ddbd.eta, frame_sfsi_dbd.eta, rel_tol=0.001), frame_sfsi_dbd.eta
    assert isclose(frame_ddbd.t_eff, frame_sfsi_dbd.t_eff, rel_tol=0.001), frame_sfsi_dbd.t_eff

def test_ddbd_frame_fixed_large():
    n_storeys = 5
    n_bays = 1
    fb = dm.FrameBuilding(n_storeys, n_bays)
    fb.material = dm.ReinforcedConcrete()
    hz = dm.Hazard()

    fb.interstorey_heights = 3.6 * np.ones(n_storeys)

    fb.bay_lengths = 6.0 * np.ones(n_bays)
    fb.set_beam_prop("depth", .6 * np.ones(n_bays))  # m      #varies vertically

    hz.z_factor = 0.3
    hz.r_factor = 1.0
    hz.n_factor = 1.0

    design_drift = 0.02

    fb.n_seismic_frames = 1
    fb.n_gravity_frames = 0

    fb.floor_length = sum(fb.bay_lengths)
    fb.floor_width = 14.0  # m

    fb.storey_masses = np.array([488.0, 488.0, 488.0, 488.0, 411.0]) * 1e3

    frame_ddbd = dbd.design_rc_frame(fb, hz, design_drift=design_drift)

    # StoreyForcesCheck1 = np.array([460402.85, 872342.25, 1235818.18, 1550830.66, 2158400.44])
    # assert isclose(frame_ddbd.v_base, 4889353.79)  # 6277794.38  # TODO: fix this
    # assert isclose(frame_ddbd.Storey_Forces, StoreyForcesCheck1)

def test_dbd_sfsi_frame_via_millen_et_al_2018():
    n_storeys = 5
    n_bays = 1
    fb = dm.FrameBuilding(n_storeys, n_bays)
    fb.n_seismic_frames = 2
    fb.n_gravity_frames = 0
    fb.material = dm.ReinforcedConcrete()
    fb.bay_lengths = [5.]
    fb.floor_width = 5.
    fb.floor_length = 5.
    fb.storey_masses = 700 * fb.floor_area * np.ones(n_storeys)
    fb.interstorey_heights = 3.4 * np.ones(n_storeys)

    fb.tie_depth = 0.8 * np.ones(fb.n_bays)  # m
    fb.tie_width = 0.8 * np.ones(fb.n_bays)  # m
    fb.foundation_rotation = 1e-3
    fb.discrete_rotation_ratio = 1.0
    fb.AxialLoadRatio = 26  # Should calculate this
    fb.Base_moment_contribution = 0.6
    fb.beam_group_size = 1
    fb.set_beam_prop('depth', 0.4, repeat='all')

    # Foundation
    fd = dm.PadFoundation()
    fd.height = 0  # m
    fd.length = 5
    fd.width = 5.
    fd.depth = 0.8
    fd.n_pads_l = 2
    fd.n_pads_w = 2

    fd.pad.width = 1.4  # m
    fd.pad.length = 1.4  # m
    fd.pad.depth = 0.5  # m
    fd.pad.height = 0
    fd.set_pad_pos_in_length_dir_as_equally_spaced()
    fd.set_pad_pos_in_width_dir_as_equally_spaced()

    fd.mass = 0
    fd2 = fd.deepcopy()

    # Soil properties

    hz = dm.Hazard()
    hz.z_factor = 0.3
    hz.r_factor = 1.0
    hz.n_factor = 1.0
    hz.corner_acc_factor = 2.
    hz.corner_period = 1

    sl = dm.Soil()
    sl.g_mod = 120e6  # Pa
    sl.poissons_ratio = 0.3  # Poisson's ratio of the soil
    sl.e_curr = 0.6  # %
    sl.phi = 35.
    sl.cohesion = 0
    sl.specific_gravity = 2.65

    design_drift = 0.02

    frame_ddbd = dbd.design_rc_frame_w_sfsi_via_millen_et_al_2018(fb, hz, sl, fd, design_drift=design_drift, verbose=2)
    assert np.isclose(frame_ddbd.delta_d, 0.08488596), frame_ddbd.delta_d
    assert np.isclose(frame_ddbd.theta_f, 0.0050136357), frame_ddbd.theta_f

def to_be_test_ddbd_sfsi_wall_from_millen_pdf_paper_2018():
    fb = dm.FrameBuilding()
    sl = dm.Soil()
    fd = dm.PadFoundation()
    hz = dm.Hazard()
    fb.id = 1
    fb.wall_depth = 3.4
    fb.wall_width = 0.3
    fb.number_walls = 4
    fb.interstorey_heights = 3.4 * np.ones(fb.n_storeys)
    fb.name = 'My' + str(fb.n_storeys) + 'storeyRBwall'  # character String
    fb.building_length = 20.0
    fb.building_width = 12.0  # m
    fb.raft_foundation = 0

    fd.height = 1.0
    fd.width = 5.5

    fb.DBaspect = 0.2
    fb.LBaspect = 2.5
    # loads:
    fb.gravity = 9.8
    fb.soil_type = 'C'
    hz.z_factor = design_spectra.calculate_z(0.4, fb.soil_type)
    hz.r_factor = 1.0
    hz.n_factor = 1.0

    fb.Live_load = 3.0e3  # Pa
    fb.floor_weight = 0.0e3  # Pa
    fb.additional_weight = 6.0e3  # Pa    #partitions, ceilings, serivces
    fb.wall_axial_contribution = 1.0
    fb.conc_weight = 23.5e3
    # Design options:

    drift_limit = 0.012
    foundation_rotation = 0.0009

    # Soil properties
    sl.g_mod = 40e6  # Pa
    sl.poissons_ratio = 0.3  # Poisson's ratio of the soil
    sl.relative_density = 0.60  # %
    sl.phi = 36
    sl.unit_moist_weight = 18e3  # N/m3
    sl.cohesion = 0

    #############################
    # Material info
    fb.fy = 300e6
    sl.E_s = 200e9
    sl.fc = 30e6  # Pa
    sl.E_conc = (3320 * np.sqrt(40.0) + 6900.0) * 1e6  # 37000000000;   #Pa    (3320*sqrt(40.0)+6900.0)*1e6
    sl.Conc_Poissons_ratio = 0.18

def load_system(n_bays=2, n_storeys=6):
    hz = sm.SeismicHazard()
    hz.z_factor = 0.3  # Hazard factor
    hz.r_factor = 1.0  # Return period factor
    hz.n_factor = 1.0  # Near-fault factor
    hz.magnitude = 7.5  # Magnitude of earthquake
    hz.corner_period = 4.0  # s
    hz.corner_acc_factor = 0.55
    sp = sm.Soil()
    sp.g_mod = 25.0e6  # [Pa]
    sp.phi = 32.0  # [degrees]
    sp.unit_dry_weight = 17000  # [N/m3]
    sp.unit_sat_weight = 18000  # [N/m3]
    sp.unit_weight_water = 9800  # [N/m3]
    sp.cohesion = 0.0  # [Pa]
    sp.poissons_ratio = 0.22

    interstorey_height = 3.4  # m
    fb = sm.FrameBuilding(n_bays=n_bays, n_storeys=n_storeys)
    fb.material = sm.materials.ReinforcedConcreteMaterial()

    fb.interstorey_heights = interstorey_height * np.ones(n_storeys)
    fb.bay_lengths = 4.0 * np.ones(n_bays)
    fb.floor_length = np.sum(fb.bay_lengths) + 4  # m
    fb.floor_width = 12.0  # m
    fb.n_seismic_frames = 3
    fb.n_gravity_frames = 0
    fb.set_storey_masses_by_pressure(9e3)  # Pa
    col_loads = fb.get_column_vert_loads()
    fb.horz2vert_mass = 1
    fb.set_beam_prop('depth', 0.6, 'all')

    fd = sm.PadFoundation()
    fd.width = fb.floor_width  # m
    fd.length = fb.floor_length  # m
    fd.depth = 0.4 + 0.1 * fb.n_storeys  # m
    fd.height = 1.0  # m
    fd.mass = 0.0  # kg
    pad = gf.size_footing_for_capacity(sp, np.max(col_loads), method='salgado', fos=3., depth=fd.depth)
    fd.pad_length = pad.length
    fd.pad_width = pad.width
    fd.pad.depth = fd.depth
    fd.pad.height = fd.height
    tie_beams_sect = sm.sections.RCBeamSection()
    tie_beams_sect.depth = fd.height
    tie_beams_sect.width = fd.height
    tie_beams_sect.rc_mat = sm.materials.ReinforcedConcreteMaterial()
    tie_beams_sect.cracked_ratio = 0.6
    fd.tie_beam_sect_in_width_dir = tie_beams_sect
    fd.tie_beam_sect_in_length_dir = tie_beams_sect
    x = fb.get_column_positions()
    x[0] = fd.pad_length / 2
    x[-1] = fd.length - fd.pad_length / 2
    fd.pad_pos_in_length_dir = x
    fd.n_pads_w = fb.n_frames
    fd.set_pad_pos_in_width_dir_as_equally_spaced()

    return fb, fd, sp, hz

def test_dbd_sfsi_frame_via_millen_et_al_2020():

    fb, fd, sp, hz = load_system(n_storeys=3, n_bays=2)
    designed_frame = dbd.design_rc_frame_w_sfsi_via_millen_et_al_2020(fb, hz, sp, fd, verbose=0)
    print('delta_ss: ', designed_frame.delta_ss)
    print('delta_f: ', designed_frame.delta_f)
    print(designed_frame.axial_load_ratio)
    assert np.isclose(designed_frame.axial_load_ratio, 4.54454554)
    assert np.isclose(designed_frame.delta_ss, 0.13432764512)
    assert np.isclose(designed_frame.delta_f, 0.0007801446), designed_frame.delta_f

def test_case_study_wall_pbd_wall_fixed_base():
    n_storeys = 6
    wb = dm.WallBuilding(n_storeys)
    wb.wall_width = 0.3  # m
    wb.wall_depth = 3.4  # m
    wb.interstorey_heights = 3.4 * np.ones(n_storeys)  # m
    wb.n_walls = 1
    wb.floor_length = 20 / 2  # m
    wb.floor_width = 12 / 2  # m
    g_load = 6000.  # Pa
    q_load = 3000.  # Pa
    eq_load_factor = 0.4
    floor_pressure = g_load + eq_load_factor * q_load
    wb.set_storey_masses_by_pressure(floor_pressure)

    # hazard
    hz = dm.Hazard()
    hz.z_factor = 0.4  # Hazard factor
    hz.r_factor = 1.0  # Return period factor
    hz.n_factor = 1.0  # Near-fault factor
    hz.magnitude = 7.5  # Magnitude of earthquake
    hz.corner_period = 3.0  # s
    hz.corner_acc_factor = 0.4

    wb.material = sm.materials.ReinforcedConcreteMaterial()
    dw = dbd.design_rc_wall(wb, hz, design_drift=0.025)

def test_case_study_wall_pbd_wall_w_sfsi():
    n_storeys = 6
    wb = dm.WallBuilding(n_storeys)
    wb.wall_width = 0.3  # m
    wb.wall_depth = 3.4  # m
    wb.interstorey_heights = 3.4 * np.ones(n_storeys)  # m
    wb.n_walls = 1
    wb.floor_length = 20 / 2  # m
    wb.floor_width = 12 / 2  # m
    g_load = 6000.  # Pa
    q_load = 3000.  # Pa
    eq_load_factor = 0.4
    floor_pressure = g_load + eq_load_factor * q_load
    wb.set_storey_masses_by_pressure(floor_pressure)

    fd = dm.RaftFoundation()
    fd.height = 1.3
    fd.length = 5.6  # m # from HDF
    fd.width = 2.25  # m # from HDF
    fd.depth = 0.0  # TODO: check this
    fd.mass = 0.0

    # soil properties from HDF
    sl = dm.Soil()
    sl.g_mod = 40e6  # Pa
    sl.poissons_ratio = 0.3
    sl.phi = 36.0  # degrees
    # sl.phi_r = np.radians(sl.phi)
    sl.cohesion = 0.0
    sl.unit_dry_weight = 18000.  # TODO: check this

    # hazard
    hz = dm.Hazard()
    hz.z_factor = 0.4  # Hazard factor
    hz.r_factor = 1.0  # Return period factor
    hz.n_factor = 1.0  # Near-fault factor
    hz.magnitude = 7.5  # Magnitude of earthquake
    hz.corner_period = 3.0  # s
    hz.corner_acc_factor = 0.4

    n_wall_eq = np.sum(wb.storey_masses) / wb.n_walls * 9.8
    n_cap_from_hdf = 12.1e6  # N
    n_wall_eq_from_hdf = 2.31e6  # N

    # n_from_input_file = (4.0e2 + 1.905e3) * 1e3
    alpha = 4.

    wb.material = sm.materials.ReinforcedConcreteMaterial()
    # dw = dbd.wall(wb, hz, design_drift=0.025)
    dw = dbd.design_rc_wall_via_millen_et_al_2020(wb, hz, sl, fd, design_drift=0.025)
    print(dw.delta_d)

def test_ddbd_wall_fixed():

    hz = dm.Hazard()
    ml.load_hazard_test_data(hz)
    wb = ml.initialise_wall_building_test_data()
    wall_dbd = dbd.design_rc_wall(wb, hz)

    assert isclose(wall_dbd.delta_d, 0.339295, rel_tol=0.001), wall_dbd.delta_d
    assert isclose(wall_dbd.mass_eff, 59429.632, rel_tol=0.001), wall_dbd.mass_eff
    assert isclose(wall_dbd.height_eff, 12.46885, rel_tol=0.001), wall_dbd.height_eff
    assert isclose(wall_dbd.mu, 1.1902, rel_tol=0.001), wall_dbd.mu
    assert isclose(wall_dbd.theta_y, 0.0168299, rel_tol=0.001), wall_dbd.theta_y
    assert isclose(wall_dbd.xi, 0.07259, rel_tol=0.001), wall_dbd.xi
    assert isclose(wall_dbd.eta, 0.86946, rel_tol=0.001), wall_dbd.eta
    assert isclose(wall_dbd.t_eff, 2.38184, rel_tol=0.001), wall_dbd.t_eff

def test_calculate_rotation_via_millen_et_al_2020():
    mom = 200.
    k_rot = 1000.0e2
    psi = 0.4
    h_eff = 3.0
    l_in = 3.0
    n_load = 300.
    n_cap = 3000.
    theta = eqdes.nonlinear_foundation.calc_fd_rot_via_millen_et_al_2020_alt_form(k_rot, l_in, n_load, n_cap, psi, mom, h_eff)
    assert np.isclose(theta, 0.0053710398), theta
    theta = eqdes.nonlinear_foundation.calc_fd_rot_via_millen_et_al_2020(k_rot, l_in, n_load, n_cap, psi, mom, h_eff)
    assert np.isclose(theta, 0.0053710398), theta

    n_load = 2000.
    n_cap = 3000.
    l_in = 5.0
    mom = 100.
    theta = eqdes.nonlinear_foundation.calc_fd_rot_via_millen_et_al_2020(k_rot, l_in, n_load, n_cap, psi, mom, h_eff)
    assert np.isclose(theta, 0.0011910855), theta

    # very large moment
    mom = 3000.
    theta = eqdes.nonlinear_foundation.calc_fd_rot_via_millen_et_al_2020(k_rot, l_in, n_load, n_cap, psi, mom, h_eff)
    assert theta is None

    # n_load equal to n_cap
    mom = 10
    n_load = 300.
    n_cap = 300.
    theta = eqdes.nonlinear_foundation.calc_fd_rot_via_millen_et_al_2020_alt_form(k_rot, l_in, n_load, n_cap, psi, mom, h_eff)
    assert theta is None

def test_calc_fd_rot_via_millen_et_al_2020_w_tie_beams():
    k_rot = 1000.0e2
    k_tbs = 100.0e2
    psi = 0.4
    h_eff = 3.0
    l_in = 3.0
    n_load = 300.
    n_cap = 3000.

    mom = 308.5

    theta_w_tbs = eqdes.nonlinear_foundation.calc_fd_rot_via_millen_et_al_2020_w_tie_beams(k_rot, l_in, n_load, n_cap, psi, mom,
                                                                                h_eff,
                                                                                k_tbs=k_tbs)

    mom_tbs = theta_w_tbs * k_tbs
    rots_adj = eqdes.nonlinear_foundation.calc_fd_rot_via_millen_et_al_2020(k_rot, l_in, n_load, n_cap, psi, mom - mom_tbs, h_eff)
    assert np.isclose(theta_w_tbs, rots_adj, rtol=0.01)

if __name__ == '__main__':
    test_dbd_sfsi_frame_via_millen_et_al_2020()
    # test_calculate_rotation_via_millen_et_al_2020()

ext: py

max_stars_repo_name: eng-tools/eqdes

max_issues_repo_name: eng-tools/eqdes

max_forks_repo_name: eng-tools/eqdes

reason: import numpy


text:
import random
import numpy as np
import copy
import pickle
from memory import Experience
from Environment import Environment
import time
import cv2
from math import isclose

class Trainer():
	def __init__(self,
				 env=None,
				 controller=None,
				 experience_memory=None,
				 quasi_newton=None,
				 batch_size=64,
				 seed=None,
				 max_iter = None,
				 **kwargs):

		self.env = env
		self.controller = controller
		self.experience_memory = experience_memory
		self.quasi_newton = quasi_newton

		# random seed
		if seed is not None:
			random.seed(seed)
			np.random.seed(seed=seed)
		self.seed = 0 if seed is None else seed

		# testing environment and parameters
		self.testing_env = Environment(task=self.env.task,seed=seed)
		self.testing_scores = [] # record testing scores
		self.epsilon_testing = 0.05
		self.max_steps_testing = 10000
		self.num_episodes_per_test = 10
		# training parameters
		self.controller_target_update_freq = 4*batch_size
		self.save_model_freq = 50000
		self.test_freq = 10000
		self.subgoal_discovery_freq = 50000
		self.epsilon_start = 1.0
		self.epsilon_end = 0.1
		self.epsilon = self.epsilon_start
		self.epsilon_annealing_steps = 1000000
		self.repeat_noop_action = 0
		self.save_results_freq = 100000
		self.batch_size = batch_size
		self.learning_starts = self.batch_size
		self.learning_freq = self.batch_size
		self.max_iter = 10000*1024 if max_iter is None else max_iter

		self.test_duration = 0.0
		self.there_was_a_test = False

		self.__dict__.update(kwargs) # updating input kwargs params

		# counters
		self.step = 0
		self.game_episode = 0

		# learning variables
		self.episode_rewards = 0.0 # including step cost
		self.episode_scores = 0.0
		# keeping records of the performance
		self.episode_rewards_list = []
		self.episode_scores_list = []
		self.episode_steps_list = []
		self.episode_time_list = []
		self.episode_start_time = 0.0
		self.episode_end_time = 0.0
		self.total_train_time = 0.0

		print('init trainer --> OK')

	def train(self):
		train_start_time = time.time()
		print('-'*60)
		print('Training DQN using L-BFGS')
		print('-'*60)
		print('-'*60)
		print('game episode: ', self.game_episode, 'time step: ', self.step)
		self.episode_start_time = time.time()
		S = self.env.reset()
		self.s = self.preprocess_concat_frames(S)
		for t in range(self.learning_starts): # fill initial expereince memory
			self.play()
			self.step += 1
		self.controller.init_Okm1() # compute g_k^{O_{k-1}} and L_k^{O_{k-1}}

		for t in range(self.max_iter+1):
			self.play()
			self.step += 1

			if (t>0) and (self.step % self.learning_freq == 0):
				self.quasi_newton.step()

			if (t>0) and (self.step % self.test_freq == 0): # test controller's performance
				self.test()

			if self.quasi_newton.termination_criterion:
				print('Quasi-Newton termination criterion --> exit')
				break

			if t>0 and (self.step % self.controller_target_update_freq == 0):
				self.controller.update_target_params()

		train_end_time = time.time()
		self.total_train_time = train_end_time - train_start_time
		self.save_results()

	def save_results(self):
		results_file_path = './results/' + 'task_' + self.env.task + \
						'_search_' + self.quasi_newton.search_method + \
						'_matrix_' + self.quasi_newton.quasi_newton_matrix + \
						'_memory_' + str(self.quasi_newton.m) + \
						'_batch_' + str(self.batch_size) + '.pkl'
		with open(results_file_path, 'wb') as f:
			pickle.dump([self.episode_steps_list,
						 self.episode_scores_list,
						 self.episode_rewards_list,
						 self.episode_time_list,
						 self.testing_scores,
						 self.total_train_time,
						 self.batch_size,
						 self.seed,
						 self.step,
						 self.quasi_newton.loss_list,
						 self.quasi_newton.grad_norm_list,
						 self.quasi_newton.computations_time_list], f)

	def save_model(self):
		model_save_path = './models/' + self.env.task + '_steps_' + str(self.step) + '.model'
		self.controller.save_model(model_save_path=model_save_path)

	def play(self):
		s = self.s
		if self.step < self.learning_starts:
			a = self.env.action_space.sample()
		else:
			a = self.epsilon_greedy()
		old_lives = self.env.lives()
		SP, r, terminal, step_info = self.env.step(a)
		new_lives = self.env.lives()
		self.episode_scores += r
		sp = self.preprocess_concat_frames(SP)

		if new_lives < old_lives:
			print('agent died, current lives = ', new_lives)
			r = min(-1.0, r)

		if (terminal and new_lives>0):
			task_done = True
			done = 1
			r = max(1.0,r)
			print('task is solved succesfully, end of episode')
		else:
			task_done = False
			done = 0

		if terminal and new_lives==0:
			print('agent terminated, end of episode')
			r = min(-1.0,r)

		if ('Pong' in self.env.task):
			if terminal and (self.episode_scores > 17):
				task_done = True
				done = 1
				r = max(1.0,r)
				print('task is solved succesfully, end of episode')
			else:
				task_done = False
				done = 0

		# if r < 0.0 or isclose(r, 0.0):
		# 	r = min(-0.01,r)

		self.episode_rewards += r

		experience = Experience(s, a, r, sp, done)
		self.experience_memory.push(experience)
		self.s = copy.deepcopy(sp)

		if terminal or task_done:
			self.episode_steps_list.append(self.step)
			self.episode_scores_list.append(self.episode_scores)
			self.episode_rewards_list.append(self.episode_rewards)
			self.episode_end_time = time.time()
			episode_time = self.episode_end_time - self.episode_start_time
			if self.there_was_a_test is True:
				episode_time = episode_time - self.test_duration
				self.there_was_a_test = False
			self.episode_time_list.append(episode_time)
			print('episode score: ', self.episode_scores)
			print('episode time: {0:.2f}' .format(episode_time))

			self.game_episode += 1
			print('-'*60)
			print('game episode: ', self.game_episode)
			print('time step: ', self.step)
			self.episode_rewards = 0.0
			self.episode_scores = 0.0
			self.episode_start_time = time.time()
			S = self.env.reset() # reset S
			self.s = self.preprocess_concat_frames(S)

	def test(self):
		self.there_was_a_test = True
		test_time_start = time.time()
		for i in range(self.num_episodes_per_test):
			self.total_score_testing = 0
			print('testing episode number: ', i)
			self.test_one_epsiode()
			self.testing_scores.append(self.total_score_testing)
			print('test score: ', self.total_score_testing)
		test_time_end = time.time()
		self.test_duration = test_time_end - test_time_start

	def test_one_epsiode(self):
		S = self.testing_env.reset()
		s = self.preprocess_concat_frames(S)
		for t in range(self.max_steps_testing):
			# self.testing_env.render()
			a = self.epsilon_greedy_testing(s)
			old_lives = self.testing_env.lives()
			SP, r, terminal, step_info = self.testing_env.step(a)
			new_lives = self.testing_env.lives()
			self.total_score_testing += r
			sp = self.preprocess_concat_frames(SP)

			if new_lives < old_lives:
				print('agent died, current lives = ', new_lives)

			if terminal and new_lives>0:
				print('task is solved succesfully, end of test episode')

			if terminal and new_lives==0:
				print('agent terminated, end of episode')

			s = copy.deepcopy(sp)

			if terminal:
				break

	def anneal_epsilon(self):
		if self.step < self.epsilon_annealing_steps:
			slop = (self.epsilon_start-self.epsilon_end)/self.epsilon_annealing_steps
			self.epsilon = self.epsilon_start - slop*self.step

	def epsilon_greedy(self):
		if random.random() < self.epsilon:
			return self.env.action_space.sample()
		else:
			return self.controller.get_best_action(self.s)

	def epsilon_greedy_testing(self,s):
		if random.random() < self.epsilon_testing:
			return self.env.action_space.sample()
		else:
			return self.controller.get_best_action(s)

	def to_grayscale(self,img):
		return np.mean(img, axis=2).astype(np.uint8)

	def downsample(self,img):
		return img[::2, ::2]

	def resize(self,img):
		w = img.shape[0]
		h = img.shape[1]
		return img.reshape((1,w,h))

	def preprocess(self,img):
		return self.resize(self.to_grayscale(self.downsample(img)))

	# def preprocess_concat_frames(self,S):
	# 	for i, img in enumerate(S):
	# 		img_preprocessed = self.preprocess(img)
	# 		if i == 0:
	# 			s = img_preprocessed
	# 		else:
	# 			s = np.concatenate((s,img_preprocessed),axis=0)
	# 	return s

	def preprocess_concat_frames(self,S):
		crop_top = 0
		crop_bottom = 0
		if 'Breakout' in self.env.task:
			crop_top = 14
			crop_bottom = 0
		if 'BeamRider' in self.env.task:
			crop_top = 20
			crop_bottom = 0
		if 'Enduro' in self.env.task:
			crop_top = 30
			crop_bottom = 40
		if 'Pong' in self.env.task:
			crop_top = 14
			crop_bottom = 5
		if 'Qbert' in self.env.task:
			crop_top = 0
			crop_bottom = 0
		if 'Seaquest' in self.env.task:
			crop_top = 20
			crop_bottom = 20
		if 'SpaceInvaders' in self.env.task:
			crop_top = 10
			crop_bottom = 6
		w = 84
		h = crop_top + 84 + crop_bottom
		for i, img in enumerate(S):
			gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
			gray_resized = cv2.resize(gray,(w,h))
			gray_cropped =  gray_resized[crop_top:h-crop_bottom,:]
			gray_reshaped = gray_cropped.reshape((1,84,84))
			if i == 0:
				s = gray_reshaped
			else:
				s = np.concatenate((s,gray_reshaped),axis=0)
		return s

max_stars_repo_name: root-master/quasi-newton-rl

max_issues_repo_name: root-master/quasi-newton-rl

max_forks_repo_name: root-master/quasi-newton-rl


text:
#!/usr/bin/env python3

# test_multiply_grad_dfx.py
#
# Copyright (c) 2010-2018 Wave Computing, Inc. and its applicable licensors.
# All rights reserved; provided, that any files identified as open source shall
# be governed by the specific open source license(s) applicable to such files.
#
# For any files associated with distributions under the Apache 2.0 license,
# full attribution to The Apache Software Foundation is given via the license
# below.
#
# PURPOSE
#      Test multiply gradients in Dynamic Fixed Point.
#
# Author          : Djordje Kovacevic
# Created On      : 05/23/2018
#
 # 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.

from collections import namedtuple

import numpy as np
import tensorflow as tf

import waveflow

def _wave_mul_grad_cc(op_x, op_y, grad):
    tuple_op = namedtuple('dummy_op', ['inputs'])
    d = tuple_op(inputs=[op_x, op_y])
    return waveflow._wave_mul_grad_dfx_cc(d, grad)

def test_mul_grad():
    ''' Run tests on the Wave custom matmul gradient operator.
    '''
    tf.reset_default_graph()

    #   .Input("grad: float32")
    #   .Input("x: float32")
    #   .Input("y: float32")
    #   .Output("grad_x: float32")
    #   .Output("grad_y: float32");
    debug = False

    t_grad = tf.get_variable("grad", [3, 2], dtype=tf.float32, initializer=tf.glorot_normal_initializer())
    t_x = tf.get_variable("x", [3, 2], dtype=tf.float32, initializer=tf.glorot_normal_initializer())
    t_y = tf.get_variable("y", [3, 2], dtype=tf.float32, initializer=tf.glorot_normal_initializer())

    t_init = tf.global_variables_initializer()

    for i in range(100):
        with tf.Session('') as sess:
            t_init.run()

            z_op = _wave_mul_grad_cc(t_x, t_y, t_grad)
            v_x, v_y, v_grad, t_z = sess.run([t_x, t_y, t_grad, z_op])
            v_z_x, v_z_y = t_z
            if debug:
                print("Iter: %d" % (i))
                print("Wave Kernel:\n-------------------------------------------------")
                print("grad: %s" % (t_grad.eval()))
                print("x: %s" % (t_x.eval()))
                print("y: %s" % (t_y.eval()))
                print("z_x: %s" % (v_z_x))
                print("z_y: %s" % (v_z_y))

            # Convert to numpy
            grad_np = np.array(v_grad)
            x_np = np.array(v_x)
            y_np = np.array(v_y)
            z_np_x = np.multiply(grad_np, y_np)
            z_np_y = np.multiply(grad_np, x_np)
            if debug:
                print("\nNumpy:\n-------------------------------------------------")
                print("grad (np): %s" % (grad_np))
                print("x (np): %s" % (x_np))
                print("y (np): %s" % (y_np))
                print("z_x (np): %s" % (z_np_x))
                print("z_y (np): %s" % (z_np_y))

            assert np.allclose(v_z_x, z_np_x, atol=1e-3)
            assert np.allclose(v_z_y, z_np_y, atol=1e-3)

    return True

if __name__ == "__main__":
    test_mul_grad()

max_stars_repo_name: wavecomp/waveext

max_issues_repo_name: kshiring/waveext

max_forks_repo_name: kshiring/waveext


text:
from .util import load_lib

import ctypes
import numpy as np
import logging
logger = logging.getLogger(__name__)

cudaLib = load_lib('libcudaDeconv')

if not cudaLib:
    logger.error('Could not load libcudaDeconv!')
else:
    try:
        # https://stackoverflow.com/questions/5862915/passing-numpy-arrays-to-a-c-function-for-input-and-output
        Deskew_interface = cudaLib.Deskew_interface
        Deskew_interface.restype = ctypes.c_int
        Deskew_interface.argtypes = [
            np.ctypeslib.ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
            ctypes.c_int,
            ctypes.c_int,
            ctypes.c_int,
            ctypes.c_float,
            ctypes.c_float,
            ctypes.c_float,
            np.ctypeslib.ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
            ctypes.c_int,
            ctypes.c_int,
            ctypes.c_float
        ]

        # Full Affine transformation
        Affine_interface = cudaLib.Affine_interface
        Affine_interface.restype = ctypes.c_int
        Affine_interface.argtypes = [
            np.ctypeslib.ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
            ctypes.c_int,
            ctypes.c_int,
            ctypes.c_int,
            np.ctypeslib.ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
            np.ctypeslib.ndpointer(ctypes.c_float, flags="C_CONTIGUOUS")
        ]

        # Affine transformation with a spatial referencing object
        Affine_interface_RA = cudaLib.Affine_interface_RA
        Affine_interface_RA.restype = ctypes.c_int
        Affine_interface_RA.argtypes = [
            np.ctypeslib.ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
            ctypes.c_int,
            ctypes.c_int,
            ctypes.c_int,
            ctypes.c_float,
            ctypes.c_float,
            ctypes.c_float,
            np.ctypeslib.ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
            np.ctypeslib.ndpointer(ctypes.c_float, flags="C_CONTIGUOUS")
        ]

    except AttributeError as e:
        logger.warning('Failed to properly import libcudaDeconv')
        print(e)

def deskewGPU(im, dxdata=0.1, dzdata=0.5, angle=31.5, width=0, shift=0, pad_val='auto'):
    """Deskew data acquired in stage-scanning mode on GPU

    Simple affine transform variant to perform a shear operation to correct
    for volume shearing.

    Args:
        im (np.ndarray): Image volume to deskew
        dxdata (float): XY Pixel size of image volume (default: {0.1})
        dzdata (float): Z-step size in image volume.  In a typical light sheet
            stage-scanning aquisition, this corresponds to the step
            size that the stage takes between planes, NOT the final
            Z-step size between planeds after deskewing along the optical
            axis of the detection objective (default: {0.5})
        angle (float): Deskew angle (usually, angle between sheet and axis of
            stage motion) (default: {31.5})
        width (int): If not 0, crop output image to specified width (default: {0})
        shift (int): If not 0, shift image center by this value (default: {0})
        pad_val (int): Value to pad image with when deskewing.  If 'auto'
            the median value of the last Z plane will be used. (default: {'auto'})

    Returns:
        np.ndarray: Deskewed volume
    """
    if isinstance(pad_val, str) and pad_val == 'auto':
        pad_val = np.median(im[-1])
    assert isinstance(pad_val, (int, float))

    nz, ny, nx = im.shape
    _dtype = im.dtype
    if not np.issubdtype(_dtype, np.float32):
        im = im.astype(np.float32)

    # have to calculate this here to know the size of the return array
    if width == 0:
        deskewedNx = np.int(nx + np.floor(nz * dzdata * abs(np.cos(angle * np.pi / 180)) / dxdata))
    else:
        deskewedNx = width

    result = np.empty((nz, ny, deskewedNx), dtype=np.float32)
    Deskew_interface(im, nx, ny, nz, dzdata, dxdata, angle, result, deskewedNx, shift, pad_val)
    return result.astype(_dtype)

def affineGPU(im, tmat, dzyx=None):
    """Perform 3D affine transformation of image given a 4x4 transformation matrix

    optional dzyx parameter {tuple, list} specifies the voxel size of the
    image [dz, dy, dx]. If it is provided, it will be used to transform
    the image from intrinsic coordinates to world coordinates prior to
    transformation, e.g.:

    *x = 0.5 + (x - 0.5) * dx*

    and then back to intrinsic coords afterwards... e.g.:

    *tu = 0.5 + (tu - 0.5) / dx*

    Args:
        im (np.ndarray): 3D input volume
        tmat (np.ndarray): Affine transformation matrix
        dzyx (list): Voxel size of input volume ([dz, dy, dx]).  If provided, the
            transformation matrix is assumed to be in units of sample space.
            otherwise the transformation is performed in image coordinates
            (default: {None})

    Returns:
        np.ndarray: Transformed volume

    Raises:
        ValueError: If the dimensions of the transformation matrix do not
            the input volume.  For instance, a 3D input volume requires a
            4 x 4 tranforamtion matrix.

    Examples:

        Perform simple translation

        >>> nx, ny, nz = (10, 20, 3)
        >>> T = np.array([[1, 0, 0, nx],
        ...               [0, 1, 0, ny],
        ...               [0, 0, 1, nz],
        ...               [0, 0, 0, 1]])
        >>> rotated = affineGPU(im, T)

        Perform a rotation about the Y axis...
        (this is the underlying code for :func:`rotateGPU`)

        >>> theta = angle * np.pi / 180
        >>> nz, ny, nx = im.shape
        >>> xzRatio = dxdata / (np.deg2rad(angle) * dzdata)
        >>> # first translate the middle of the image to the origin
        >>> T1 = np.array([[1, 0, 0, nx / 2],
        >>>                [0, 1, 0, ny / 2],
        >>>                [0, 0, 1, nz / 2],
        >>>                [0, 0, 0, 1]])
        >>> # then scale (resample) the Z axis the dz/dx ratio
        >>> S = np.array([[1, 0, 0, 0],
        >>>               [0, 1, 0, 0],
        >>>               [0, 0, xzRatio, 0],
        >>>               [0, 0, 0, 1]])
        >>> # then rotate theta degrees about the Y axis
        >>> R = np.array([[np.cos(theta), 0, -np.sin(theta), 0],
        >>>               [0, 1, 0, 0],
        >>>               [np.sin(theta), 0, np.cos(theta), 0],
        >>>               [0, 0, 0, 1]])
        >>> # then translate back to the original origin
        >>> T2 = np.array([[1, 0, 0, -nx / 2],
        >>>                [0, 1, 0, -ny / 2],
        >>>                [0, 0, 1, -nz / 2],
        >>>                [0, 0, 0, 1]])
        >>> T = np.eye(4)
        >>> T = np.dot(np.dot(np.dot(np.dot(T, T1), S), R), T2)
        >>> rotated = affineGPU(im, T)

    """
    if not tmat.shape == tuple([im.ndim + 1] * 2):
        raise ValueError('{} dimensional transform matrix used on {} dimensional image'
                         .format(tmat.shape[0] - 1, im.ndim))

    nz, ny, nx = im.shape
    if not np.issubdtype(im.dtype, np.float32) or not im.flags['C_CONTIGUOUS']:
        im = np.ascontiguousarray(im, dtype=np.float32)
    if not np.issubdtype(tmat.dtype, np.float32):
        tmat = tmat.astype(np.float32)
    # have to calculate this here to know the size of the return array
    result = np.empty((nz, ny, nx), dtype=np.float32)
    if (isinstance(dzyx, (tuple, list)) and
            len(dzyx) == 3):
        dzyx = [float(i) for i in dzyx]
        # note, dzyx coordinate order is flipped when handing to Affine_interface_RA
        Affine_interface_RA(im, nx, ny, nz, dzyx[2], dzyx[1], dzyx[0], result, tmat)
    else:
        Affine_interface(im, nx, ny, nz, result, tmat)
    return result

def rotateGPU(im, dzdata, dxdata=0.1, angle=31.5, reverse=False):
    """Rotate image around Y axis by some angle

    This is a convenience function that will apply the appropriate affine
    transformation for rotating a volume around the Y axis by some angle.
    This is typically done with images acquired on inverted light sheet
    microscopes where the image plane is not parallel to the coverslip
    (such as lattice light sheet, or diSPIM microscopes), in order to change
    the coordinate system of the image volume such that the Z axis is
    orthogonal to the coverslip

    Args:
        im (np.ndarray): 3D volume to be rotated
        dzdata (float): Z-step size in microns of the image volume
        dxdata (float): XY pixel size of the volume (default: {0.1})
        angle (float): Angle to rotate around Y axis (default: {31.5})
        reverse (bool): Rotate in the opposite direction. (default: {False})

    Returns:
        np.ndarray: Rotated 3D volume
    """
    angle = float(angle)
    xzRatio = dxdata / (np.deg2rad(angle) * dzdata)

    npad = ((0, 0), (0, 0), (0, 0))
    im = np.pad(im, pad_width=npad, mode='constant', constant_values=0)

    theta = angle * np.pi / 180
    theta = theta if not reverse else -theta

    nz, ny, nx = im.shape
    # first translate the middle of the image to the origin
    T1 = np.array([[1, 0, 0, nx / 2],
                   [0, 1, 0, ny / 2],
                   [0, 0, 1, nz / 2],
                   [0, 0, 0, 1]])
    # then scale (resample) the Z axis the dz/dx ratio
    S = np.array([[1, 0, 0, 0],
                  [0, 1, 0, 0],
                  [0, 0, xzRatio, 0],
                  [0, 0, 0, 1]])
    # then rotate theta degrees about the Y axis
    R = np.array([[np.cos(theta), 0, -np.sin(theta), 0],
                  [0, 1, 0, 0],
                  [np.sin(theta), 0, np.cos(theta), 0],
                  [0, 0, 0, 1]])
    # then translate back to the original origin
    T2 = np.array([[1, 0, 0, -nx / 2],
                   [0, 1, 0, -ny / 2],
                   [0, 0, 1, -nz / 2],
                   [0, 0, 0, 1]])
    T = np.eye(4)
    T = np.dot(np.dot(np.dot(np.dot(T, T1), S), R), T2)

    rotated = affineGPU(im, T)

    return rotated

max_stars_repo_name: Jun-NIBS/pycudadecon

max_issues_repo_name: Jun-NIBS/pycudadecon

max_forks_repo_name: Jun-NIBS/pycudadecon


text:
# This file is part of GridCal.
#
# GridCal is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# GridCal is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with GridCal.  If not, see <http://www.gnu.org/licenses/>.

import os
import time

import numpy as np
import numba as nb
from scipy.sparse import lil_matrix, diags, csc_matrix

import GridCal.Engine.Simulations.PowerFlow.derivatives as deriv

@nb.njit()
def make_lookup(n, arr):
    lookup = np.zeros(n, dtype=np.int32)
    lookup[arr] = np.arange(len(arr), dtype=np.int32)
    return lookup

@nb.njit()
def fill_acdc_jacobian_data(Jx, Ji, Jp, Yp, Yi, Ys,
                            dSbus_dVa_x, dSbus_dVm_x,
                            dSf_dVa_x, dSf_dVa_i, dSf_dVa_p,
                            dSf_dVm_x, dSf_dVm_i, dSf_dVm_p,
                            dSt_dVa_x, dSt_dVa_i, dSt_dVa_p,
                            dSt_dVm_x, dSt_dVm_i, dSt_dVm_p,
                            dPfdp_dVm_x, dPfdp_dVm_i, dPfdp_dVm_p,
                            pvpq, pvpq_lookup, npvpq, pq,
                            i2, ni2, i2_lookup, i4, ni4, i4_lookup,
                            iPfsh, nPfsh, iPfsh_lookup,
                            iQtma, nQtma, iQtma_lookup,
                            iQfma, nQfma,
                            iVtma, nVtma,
                            iPfdp, nPfdp, iPfdp_lookup,
                            iBeqz, nBeqz,
                            iBeqv, nBeqv,
                            F, T, V, ma, k2, tap, Bc, Beq):
    """

    :param Jx:
    :param Ji:
    :param Jp:
    :param Yp:
    :param Yi:
    :param Ys:
    :param dSbus_dVa_x:
    :param dSbus_dVm_x:
    :param dSf_dVa_x:
    :param dSf_dVa_i:
    :param dSf_dVa_p:
    :param dSf_dVm_x:
    :param dSf_dVm_i:
    :param dSf_dVm_p:
    :param dSt_dVa_x:
    :param dSt_dVa_i:
    :param dSt_dVa_p:
    :param dSt_dVm_x:
    :param dSt_dVm_i:
    :param dSt_dVm_p:
    :param dPfdp_dVm_x:
    :param dPfdp_dVm_i:
    :param dPfdp_dVm_p:
    :param pvpq: Array of pv and then pq bus indices (not sorted)
    :param pvpq_lookup:
    :param npvpq:
    :param pvpq: Array of pq bus indices (sorted)
    :param i2:
    :param ni2:
    :param i2_lookup:
    :param i4:
    :param ni4:
    :param i4_lookup:
    :param iPfsh: indices of the Pf controlled branches
    :param nPfsh:
    :param iPfsh_lookup:
    :param iQtma: Indices of the Qt controlled branches
    :param nQtma:
    :param iQtma_lookup:
    :param iQfma: indices of the Qf controlled branches
    :param nQfma:
    :param iVtma: Indices of the Vt controlled branches
    :param nVtma:
    :param iPfdp: indices of the droop controlled branches
    :param nPfdp:
    :param iPfdp_lookup:
    :param iBeqz: Indices of the Qf controlled branches
    :param nBeqz:
    :param iBeqv: Indices of the Vf Controlled branches
    :param nBeqv:
    :param F: Array of "from" bus indices
    :param T: Array of "to" bus indices
    :param V: Array of complex bus voltages
    :param ma: Array of tap modules
    :param k2: Array of branch converter losses
    :param tap: Array of complex tap values {remember tap = ma * exp(1j * theta) }
    :param Bc: Array of branch full susceptances
    :param Beq: Array of branch equivalent (variable) susceptances
    :return:
    """

    nnz = 0
    p = 0
    Jp[p] = 0

    # column 1: derivatives w.r.t Va -----------------------------------------------------------------------------------
    for j in pvpq:  # sliced columns

        # J11
        if npvpq:
            for k in range(Yp[j], Yp[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = Yi[k]
                ii = pvpq_lookup[i]

                if pvpq[ii] == i:
                    # entry found
                    Jx[nnz] = dSbus_dVa_x[k].real
                    Ji[nnz] = ii
                    nnz += 1

        # J21 J31 J41
        offset = npvpq
        if ni2:
            for k in range(Yp[j], Yp[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = Yi[k]
                ii = i2_lookup[i]

                if i2[ii] == i:
                    # entry found
                    Jx[nnz] = dSbus_dVa_x[k].imag
                    Ji[nnz] = ii + offset
                    nnz += 1

        # J51
        offset += ni2
        if nPfsh:
            for k in range(dSf_dVa_p[j], dSf_dVa_p[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = dSf_dVa_i[k]
                ii = iPfsh_lookup[i]

                if iPfsh[ii] == i:
                    # entry found
                    Jx[nnz] = dSf_dVa_x[k].real
                    Ji[nnz] = ii + offset
                    nnz += 1

        # J61 J71
        offset += nPfsh
        if ni4:
            for k in range(dSf_dVa_p[j], dSf_dVa_p[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = dSf_dVa_i[k]
                ii = i4_lookup[i]

                if i4[ii] == i:
                    # entry found
                    Jx[nnz] = dSf_dVa_x[k].imag
                    Ji[nnz] = ii + offset
                    nnz += 1

        # J81
        offset += ni4
        if nQtma:
            for k in range(dSt_dVa_p[j], dSt_dVa_p[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = dSt_dVa_i[k]
                ii = iQtma_lookup[i]

                if iQtma[ii] == i:
                    # entry found
                    Jx[nnz] = dSt_dVa_x[k].imag
                    Ji[nnz] = ii + offset
                    nnz += 1

        # J91
        offset += nQtma
        if nPfdp:
            for k in range(dSf_dVa_p[j], dSf_dVa_p[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = dSf_dVa_i[k]
                ii = iPfdp_lookup[i]

                if iPfdp[ii] == i:
                    # entry found
                    Jx[nnz] = -dSf_dVa_x[k].real
                    Ji[nnz] = ii + offset
                    nnz += 1

        # finalize column
        p += 1
        Jp[p] = nnz

    # column 2: derivatives w.r.t Vm -----------------------------------------------------------------------------------
    for j in pq:  # sliced columns

        # J11
        if npvpq:
            for k in range(Yp[j], Yp[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = Yi[k]
                ii = pvpq_lookup[i]

                if pvpq[ii] == i:
                    # entry found
                    Jx[nnz] = dSbus_dVm_x[k].real
                    Ji[nnz] = ii
                    nnz += 1

        # J21 J31 J41
        offset = npvpq
        if ni2:
            for k in range(Yp[j], Yp[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = Yi[k]
                ii = i2_lookup[i]

                if i2[ii] == i:
                    # entry found
                    Jx[nnz] = dSbus_dVm_x[k].imag
                    Ji[nnz] = ii + offset
                    nnz += 1

        # J51
        offset += ni2
        if nPfsh:
            for k in range(dSf_dVm_p[j], dSf_dVm_p[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = dSf_dVm_i[k]
                ii = iPfsh_lookup[i]

                if iPfsh[ii] == i:
                    # entry found
                    Jx[nnz] = dSf_dVm_x[k].real
                    Ji[nnz] = ii + offset
                    nnz += 1

        # J61 J71
        offset += nPfsh
        if ni4:
            for k in range(dSf_dVm_p[j], dSf_dVm_p[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = dSf_dVm_i[k]
                ii = i4_lookup[i]

                if i4[ii] == i:
                    # entry found
                    Jx[nnz] = dSf_dVm_x[k].imag
                    Ji[nnz] = ii + offset
                    nnz += 1

        # J81
        offset += ni4
        if nQtma:
            for k in range(dSt_dVm_p[j], dSt_dVm_p[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = dSt_dVm_i[k]
                ii = iQtma_lookup[i]

                if iQtma[ii] == i:
                    # entry found
                    Jx[nnz] = dSt_dVm_x[k].imag
                    Ji[nnz] = ii + offset
                    nnz += 1

        # J91
        offset += nQtma
        if nPfdp:

            for k in range(dPfdp_dVm_p[j], dPfdp_dVm_p[j + 1]):  # rows of A[:, j]

                # row index translation to the "rows" space
                i = dPfdp_dVm_i[k]
                ii = iPfdp_lookup[i]

                if iPfdp[ii] == i:
                    # entry found
                    Jx[nnz] = dPfdp_dVm_x[k]
                    Ji[nnz] = ii + offset
                    nnz += 1

        # finalize column
        p += 1
        Jp[p] = nnz

    # Column 3: derivatives w.r.t Beq for iBeqz + iBeqv ----------------------------------------------------------------
    if nBeqz + nBeqv:
        indices = np.concatenate((iBeqz, iBeqv))
        dSbus_dBeq_data, dSbus_dBeq_indices, dSbus_dBeq_indptr, \
        dSf_dBeqx_data, dSf_dBeqx_indices, dSf_dBeqx_indptr = deriv.derivatives_Beq_csc_numba(indices,
                                                                                F, V, ma, k2)

        for j in range(nBeqz + nBeqv):  # sliced columns

            # J11
            if npvpq:
                for k in range(dSbus_dBeq_indptr[j], dSbus_dBeq_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSbus_dBeq_indices[k]
                    ii = pvpq_lookup[i]

                    if pvpq[ii] == i:
                        # entry found
                        Jx[nnz] = dSbus_dBeq_data[k].real
                        Ji[nnz] = ii
                        nnz += 1

            # J21 J31 J41
            offset = npvpq
            if ni2:
                for k in range(dSbus_dBeq_indptr[j], dSbus_dBeq_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSbus_dBeq_indices[k]
                    ii = i2_lookup[i]

                    if i2[ii] == i:
                        # entry found
                        Jx[nnz] = dSbus_dBeq_data[k].imag
                        Ji[nnz] = ii + offset
                        nnz += 1

            # J51
            offset += ni2
            if nPfsh:
                for k in range(dSf_dBeqx_indptr[j], dSf_dBeqx_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSf_dBeqx_indices[k]
                    ii = iPfsh_lookup[i]

                    if iPfsh[ii] == i:
                        # entry found
                        Jx[nnz] = dSf_dBeqx_data[k].real
                        Ji[nnz] = ii + offset
                        nnz += 1

            # J61 J71
            offset += nPfsh
            if ni4:
                for k in range(dSf_dBeqx_indptr[j], dSf_dBeqx_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSf_dBeqx_indices[k]
                    ii = i4_lookup[i]

                    if i4[ii] == i:
                        # entry found
                        Jx[nnz] = dSf_dBeqx_data[k].imag
                        Ji[nnz] = ii + offset
                        nnz += 1

            # J81
            offset += ni4
            # if nQtma:  # --> The Jacobian is always zero :|
            #     for k in range(dSt_dBeqz.indptr[j], dSt_dBeqz.indptr[j + 1]):  # rows of A[:, j]
            #
            #         # row index translation to the "rows" space
            #         i = dSt_dBeqz.indices[k]
            #         ii = iQtma_lookup[i]
            #
            #         if iQtma[ii] == i:
            #             # entry found
            #             Jx[nnz] = dSt_dBeqz.data[k].imag
            #             Ji[nnz] = ii + offset
            #             nnz += 1

            # J91
            offset += nQtma
            if nPfdp:
                for k in range(dSf_dBeqx_indptr[j], dSf_dBeqx_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSf_dBeqx_indices[k]
                    ii = iPfdp_lookup[i]

                    if iPfdp[ii] == i:
                        # entry found
                        Jx[nnz] = -dSf_dBeqx_data[k].real
                        Ji[nnz] = ii + offset
                        nnz += 1

            # finalize column
            p += 1
            Jp[p] = nnz

    # Column 4: derivative w.r.t ma for iQfma + iQfma + iVtma ----------------------------------------------------------
    if nQfma + nQtma + nVtma:
        indices = np.concatenate((iQfma, iQtma, iVtma))
        dSbus_dma_data, dSbus_dma_indices, dSbus_dma_indptr, \
        dSf_dma_data, dSf_dma_indices, dSf_dma_indptr, \
        dSt_dma_data, dSt_dma_indices, dSt_dma_indptr = deriv.derivatives_ma_csc_numba(indices,
                                                                                F, T, Ys, k2, tap, ma, Bc,
                                                                                Beq, V)

        for j in range(nQfma + nQtma + nVtma):  # sliced columns

            # J11
            if npvpq:
                for k in range(dSbus_dma_indptr[j], dSbus_dma_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSbus_dma_indices[k]
                    ii = pvpq_lookup[i]

                    if pvpq[ii] == i:
                        # entry found
                        Jx[nnz] = dSbus_dma_data[k].real
                        Ji[nnz] = ii
                        nnz += 1

            # J21 J31 J41
            offset = npvpq
            if ni2:
                for k in range(dSbus_dma_indptr[j], dSbus_dma_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSbus_dma_indices[k]
                    ii = i2_lookup[i]

                    if i2[ii] == i:
                        # entry found
                        Jx[nnz] = dSbus_dma_data[k].imag
                        Ji[nnz] = ii + offset
                        nnz += 1

            # J51
            offset += ni2
            if nPfsh:
                for k in range(dSf_dma_indptr[j], dSf_dma_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSf_dma_indices[k]
                    ii = iPfsh_lookup[i]

                    if iPfsh[ii] == i:
                        # entry found
                        Jx[nnz] = dSf_dma_data[k].real
                        Ji[nnz] = ii + offset
                        nnz += 1

            # J61 J71
            offset += nPfsh
            if ni4:
                for k in range(dSf_dma_indptr[j], dSf_dma_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSf_dma_indices[k]
                    ii = i4_lookup[i]

                    if i4[ii] == i:
                        # entry found
                        Jx[nnz] = dSf_dma_data[k].imag
                        Ji[nnz] = ii + offset
                        nnz += 1

            # J81
            offset += ni4
            if nQtma:
                for k in range(dSt_dma_indptr[j], dSt_dma_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSt_dma_indices[k]
                    ii = iQtma_lookup[i]

                    if iQtma[ii] == i:
                        # entry found
                        Jx[nnz] = dSt_dma_data[k].imag
                        Ji[nnz] = ii + offset
                        nnz += 1

            # J91
            offset += nQtma
            if nPfdp:
                for k in range(dSf_dma_indptr[j], dSf_dma_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSf_dma_indices[k]
                    ii = iPfdp_lookup[i]

                    if iPfdp[ii] == i:
                        # entry found
                        Jx[nnz] = -dSf_dma_data[k].real
                        Ji[nnz] = ii + offset
                        nnz += 1

            # finalize column
            p += 1
            Jp[p] = nnz

    # Column 5: derivatives w.r.t theta sh for iPfsh + droop -----------------------------------------------------------
    if nPfsh + nPfdp > 0:

        indices = np.concatenate((iPfsh, iPfdp))
        dSbus_dsh_data, dSbus_dsh_indices, dSbus_dsh_indptr, \
            dSf_dsh_data, dSf_dsh_indices, dSf_dsh_indptr, \
            dSt_dsh_data, dSt_dsh_indices, dSt_dsh_indptr = deriv.derivatives_sh_csc_numba(indices,
                                                                                F, T, Ys, k2, tap, V)

        for j in range(nPfsh + nPfdp):  # sliced columns

            # J11
            if npvpq:
                for k in range(dSbus_dsh_indptr[j], dSbus_dsh_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSbus_dsh_indices[k]
                    ii = pvpq_lookup[i]

                    if pvpq[ii] == i:
                        # entry found
                        Jx[nnz] = dSbus_dsh_data[k].real
                        Ji[nnz] = ii
                        nnz += 1

            # J21 J31 J41
            offset = npvpq
            if ni2:
                for k in range(dSbus_dsh_indptr[j], dSbus_dsh_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSbus_dsh_indices[k]
                    ii = i2_lookup[i]

                    if i2[ii] == i:
                        # entry found
                        Jx[nnz] = dSbus_dsh_data[k].imag
                        Ji[nnz] = ii + offset
                        nnz += 1

            # J51
            offset += ni2
            if nPfsh:
                for k in range(dSf_dsh_indptr[j], dSf_dsh_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSf_dsh_indices[k]
                    ii = iPfsh_lookup[i]

                    if iPfsh[ii] == i:
                        # entry found
                        Jx[nnz] = dSf_dsh_data[k].real
                        Ji[nnz] = ii + offset
                        nnz += 1

            # J61 J71
            offset += nPfsh
            if ni4:
                for k in range(dSf_dsh_indptr[j], dSf_dsh_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSf_dsh_indices[k]
                    ii = i4_lookup[i]

                    if i4[ii] == i:
                        # entry found
                        Jx[nnz] = dSf_dsh_data[k].imag
                        Ji[nnz] = ii + offset
                        nnz += 1

            # J81
            offset += ni4
            if nQtma:
                for k in range(dSt_dsh_indptr[j], dSt_dsh_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSt_dsh_indices[k]
                    ii = iQtma_lookup[i]

                    if iQtma[ii] == i:
                        # entry found
                        Jx[nnz] = dSt_dsh_data[k].imag
                        Ji[nnz] = ii + offset
                        nnz += 1

            # J91
            offset += nQtma
            if nPfdp:
                for k in range(dSf_dsh_indptr[j], dSf_dsh_indptr[j + 1]):  # rows of A[:, j]

                    # row index translation to the "rows" space
                    i = dSf_dsh_indices[k]
                    ii = iPfdp_lookup[i]

                    if iPfdp[ii] == i:
                        # entry found
                        Jx[nnz] = -dSf_dsh_data[k].real
                        Ji[nnz] = ii + offset
                        nnz += 1

            # finalize column
            p += 1
            Jp[p] = nnz

    # Finalize ----------------------------------------------------------------------------
    #  finalize the Jacobian Pointer
    Jp[p] = nnz

    return nnz, Jx, Ji, Jp

def fubm_jacobian(nb, nl, iPfsh, iPfdp, iQfma, iQtma, iVtma, iBeqz, iBeqv, VfBeqbus, Vtmabus,
                  F, T, Ys, k2, tap, ma, Bc, Beq, Kdp, V, Ybus, Yf, Yt, Cf, Ct, pvpq, pq):
    """
    Compute the FUBM jacobian in a dynamic fashion by only computing the derivatives that are needed
    :param nb: number of buses
    :param nl: Number of lines
    :param iPfsh: indices of the Pf controlled branches
    :param iPfdp: indices of the droop controlled branches
    :param iQfma: indices of the Qf controlled branches
    :param iQtma: Indices of the Qt controlled branches
    :param iVtma: Indices of the Vt controlled branches
    :param iBeqz: Indices of the Qf controlled branches
    :param iBeqv: Indices of the Vf Controlled branches
    :param F: Array of "from" bus indices
    :param T: Array of "to" bus indices
    :param Ys: Array of branch series admittances
    :param k2: Array of branch converter losses
    :param tap: Array of complex tap values {remember tap = ma * exp(1j * theta) }
    :param ma: Array of tap modules
    :param Bc: Array of branch full susceptances
    :param Beq: Array of branch equivalent (variable) susceptances
    :param Kdp: Array of branch converter droop constants
    :param V: Array of complex bus voltages
    :param Ybus: Admittance matrix
    :param Yf: Admittances matrix of the branches with the "from" buses
    :param Yt: Admittances matrix of the branches with the "to" buses
    :param Cf: Connectivity matrix of the branches with the "from" buses
    :param Ct: Connectivity matrix of the branches with the "to" buses
    :param pvpq: Array of pv and then pq bus indices (not sorted)
    :param pq: Array of PQ bus indices
    :return: FUBM Jacobian matrix
    """
    nPfsh = len(iPfsh)
    nPfdp = len(iPfdp)
    nQfma = len(iQfma)
    nQtma = len(iQtma)
    nVtma = len(iVtma)
    nBeqz = len(iBeqz)
    nBeqv = len(iBeqv)
    nVfBeqbus = len(VfBeqbus)
    nVtmabus = len(Vtmabus)
    npq = len(pq)
    npvpq = len(pvpq)
    nbus = Ybus.shape[0]
    nbr = Yf.shape[0]

    i2 = np.r_[pq, VfBeqbus, Vtmabus]
    i4 = np.r_[iQfma, iBeqz]
    ni2 = len(i2)
    ni4 = len(i4)

    # compose the derivatives of the power injections w.r.t Va and Vm
    dSbus_dVm_x, dSbus_dVa_x = deriv.dSbus_dV_numba_sparse_csc(Ybus.data, Ybus.indptr, Ybus.indices, V)

    # compose the derivatives of the branch flow w.r.t Va and Vm
    Vc = np.conj(V)
    E = V / np.abs(V)
    dSf_dVa, dSf_dVm = deriv.dSf_dV_fast(Yf, V, Vc, E, F, Cf)

    if nQtma:
        dSt_dVa, dSt_dVm = deriv.dSf_dV_fast(Yt, V, Vc, E, T, Ct)
    else:
        dSt_dVa = csc_matrix((nl, nb))
        dSt_dVm = csc_matrix((nl, nb))

    if nPfdp:
        # compute the droop derivative
        dVmf_dVm = lil_matrix((nl, nb))
        dVmf_dVm[iPfdp, :] = Cf[iPfdp, :]
        dPfdp_dVm = -dSf_dVm.real + diags(Kdp) * dVmf_dVm
    else:
        dPfdp_dVm = csc_matrix((nPfdp, nb))

    n_cols = npvpq + npq + nBeqz + nBeqv + nQfma + nQtma + nVtma + nPfsh + nPfdp
    n_rows = n_cols

    nnz_estimate = Ybus.nnz * 8
    Jx = np.empty(nnz_estimate, dtype=np.float64)  # data
    Ji = np.empty(nnz_estimate, dtype=np.int32)  # indices
    Jp = np.empty(n_cols + 1, dtype=np.int32)  # pointers

    # generate lookup for the row slicing
    pvpq_lookup = make_lookup(nbus, pvpq)
    i2_lookup = make_lookup(nbus, i2)
    iPfsh_lookup = make_lookup(nbr, iPfsh)
    i4_lookup = make_lookup(nbr, i4)
    iQtma_lookup = make_lookup(nbr, iQtma)
    iPfdp_lookup = make_lookup(nbr, iPfdp)

    # fill the jacobian data with numba
    nnz, Jx, Ji, Jp = fill_acdc_jacobian_data(Jx, Ji, Jp, Ybus.indptr, Ybus.indices, Ys,
                                              dSbus_dVa_x, dSbus_dVm_x,
                                              dSf_dVa.data, dSf_dVa.indices, dSf_dVa.indptr,
                                              dSf_dVm.data, dSf_dVm.indices, dSf_dVm.indptr,
                                              dSt_dVa.data, dSt_dVa.indices, dSt_dVa.indptr,
                                              dSt_dVm.data, dSt_dVm.indices, dSt_dVm.indptr,
                                              dPfdp_dVm.data, dPfdp_dVm.indices, dPfdp_dVm.indptr,
                                              pvpq, pvpq_lookup, npvpq, pq,
                                              i2, ni2, i2_lookup, i4, ni4, i4_lookup,
                                              iPfsh, nPfsh, iPfsh_lookup,
                                              iQtma, nQtma, iQtma_lookup,
                                              iQfma, nQfma,
                                              iVtma, nVtma,
                                              iPfdp, nPfdp, iPfdp_lookup,
                                              iBeqz, nBeqz,
                                              iBeqv, nBeqv,
                                              F, T, V, ma, k2, tap, Bc, Beq)

    Jx = np.resize(Jx, nnz)
    Ji = np.resize(Ji, nnz)

    J = csc_matrix((Jx, Ji, Jp), shape=(n_rows, n_cols))

    return J

reason: import numpy,from scipy,import numba


text:
import numpy as np
import pandas as pd
import spacy
import random
from copy import deepcopy

from checklist.perturb import Perturb
from allennlp.data.tokenizers.spacy_tokenizer import SpacyTokenizer

from src.data.constants import \
    ADJECTIVES, \
    REVERSE_CONTRACTION_MAP, \
    CONTRACTION_MAP, \
    PUNCTUATION, \
    DICT_GENDER

nlp = spacy.load('en_core_web_sm')

def _custom_remove_char(text_orig, char):
    """Removes characters from a list of string
    Inputs:
    text: String or list that has strings as elements. Transformation will be applied to all strings.
    char: character(s) to be removed. This can be a string, or a list of strings (if multiple characters
    need to be removed)

    Output: modified string, or list of modified strings
    """
    if type(text_orig) == str:
        text = [text_orig]
    else:
        text = text_orig
    if type(char) == str:
        char = [char]
    result = []
    dummy_string = ''
    dummy_dict = dict.fromkeys(char, '')
    table = dummy_string.maketrans(dummy_dict)
    for i in range(len(text)):
        result.append(text[i].translate(table))
    if type(text_orig) == str:
        result = result[0]
    return result

def _gen_empty_columns():
    new_column_tokens = []
    new_column_concat = []
    new_column_success_flag = []
    new_column_pert_indices = []
    return new_column_tokens, new_column_concat, new_column_success_flag, new_column_pert_indices

def swap_adjectives(df: pd.DataFrame, orig_index: int, tokens_orig: list):
    """
    swap two consecutive adjectives (connected by 'and' or 'or')

    :param df: DataFrame containing sentences
    :param orig_index: integer indicating the column index of the original sentence
    :param tokens_orig: tokenised version of sentence
    :return: None. Modifies DataFrame in-place
    """

    new_column_tokens, new_column_concat, new_column_success_flag, new_column_pert_indices = _gen_empty_columns()

    for s in range(len(tokens_orig)):
        sentence = deepcopy(tokens_orig[s])
        pert_indices = []
        for t in range(1, len(sentence)-1):
            if sentence[t] in ['and', 'or']:
                if sentence[t - 1] in ADJECTIVES and sentence[t + 1] in ADJECTIVES:
                    adj_1 = str(sentence[t - 1])
                    adj_2 = str(sentence[t + 1])
                    sentence[t - 1] = adj_2
                    sentence[t + 1] = adj_1
                    pert_indices.extend([t - 1, t + 1])
        new_column_tokens.append(sentence)
        if len(pert_indices) == 0:
            new_column_concat.append(df.iloc[s, orig_index])
            new_column_success_flag.append(0)
            new_column_pert_indices.append(None)
        else:
            new_column_concat.append(" ".join(sentence))
            new_column_success_flag.append(1)
            new_column_pert_indices.append(pert_indices)

    df['swap_adj_concat'] = new_column_concat
    df['swap_adj_tokens'] = new_column_tokens
    df['swap_adj_success'] = new_column_success_flag
    df['swap_adj_pert_ind'] = new_column_pert_indices

def contraction(df: pd.DataFrame, orig_index: int, tokens_orig: list):
    """
    Change to and from contracted form (e.g. "you're" to "you are", or "you are" to "you're")

    :param df: DataFrame containing sentences
    :param orig_index: integer indicating the column index of the original sentence
    :param tokens_orig: tokenised version of sentence
    :return: None. Modifies DataFrame in-place
    """

    new_column_tokens, new_column_concat, new_column_success_flag, new_column_pert_indices = _gen_empty_columns()

    for s in range(len(tokens_orig)):
        sentence = deepcopy(tokens_orig[s])
        pert_indices = []
        t = 0
        while t < len(sentence) - 2:
            n_grams_with_space = {}
            n_grams_no_space = {}
            change_flag = 0
            for i in range(1, 3 + 1):
                n_grams_no_space[i] = "".join(sentence[t:t + i])
                n_grams_with_space[i] = " ".join(sentence[t:t + i])
            for n in range(1, 3 + 1):
                n_gram_no_space = n_grams_no_space[n]
                n_gram_with_space = n_grams_with_space[n]
                if n_gram_no_space in REVERSE_CONTRACTION_MAP:
                    new_phrase = REVERSE_CONTRACTION_MAP[n_gram_no_space]
                    # if number of tokens before and after modification is not the same, leave sentence unchanged
                    if len(new_phrase) != n:
                        continue
                    else:
                        for i in range(n):
                            sentence[t + i] = new_phrase[i]
                    for i in range(n):
                        pert_indices.append(t + i)
                    t += n
                    change_flag += 1
                    break
                elif n_gram_with_space in CONTRACTION_MAP:
                    new_phrase = CONTRACTION_MAP[n_gram_with_space]
                    # if number of tokens before and after modification is not the same, leave sentence unchanged
                    if len(new_phrase) != n:
                        continue
                    else:
                        for i in range(n):
                            sentence[t + i] = new_phrase[i]
                    for i in range(n):
                        pert_indices.append(t + i)
                    t += n
                    change_flag += 1
                    break
            # if no change has been applied to any of the n-grams, move to next token in sentence
            if change_flag == 0:
                t += 1

        if len(pert_indices) == 0:
            new_column_tokens.append(sentence)
            new_column_concat.append(df.iloc[s, orig_index])
            new_column_success_flag.append(0)
            new_column_pert_indices.append(None)
        else:
            new_column_tokens.append(sentence)
            new_column_concat.append(" ".join(sentence))
            new_column_success_flag.append(1)
            new_column_pert_indices.append(pert_indices)

    df['contraction_concat'] = new_column_concat
    df['contraction_tokens'] = new_column_tokens
    df['contraction_success'] = new_column_success_flag
    df['contraction_pert_ind'] = new_column_pert_indices

def change_first_name(df: pd.DataFrame, orig_index: int, tokens_orig: list):
    """
    Change first name in sentence if one exists and if name is in CheckList's name lookup json

    :param df: DataFrame containing sentences
    :param orig_index: integer indicating the column index of the original sentence
    :param tokens_orig: tokenised version of sentence
    :return: None. Modifies DataFrame in-place
    """

    new_column_tokens, new_column_concat, new_column_success_flag, new_column_pert_indices = _gen_empty_columns()
    # Checklist requires pre-processing with Spacy for this perturbation
    pdata = list(nlp.pipe(df.iloc[:, orig_index]))

    for s in range(len(tokens_orig)):
        sentence = tokens_orig[s]
        new_sentence = Perturb.change_names(pdata[s], n=1, first_only=True, meta=True)
        if not new_sentence:
            new_column_tokens.append(sentence)
            new_column_concat.append(df.iloc[s, orig_index])
            new_column_success_flag.append(0)
            new_column_pert_indices.append(None)
        else:
            # extract token that has been perturbed
            token_pert = new_sentence[1][0][0]
            # verify that perturbed name appears as token in the original input
            if token_pert in sentence:
                # obtain index
                token_index = sentence.index(token_pert)
                new_sentence_tokens = deepcopy(sentence)
                # replace token
                new_sentence_tokens[token_index] = new_sentence[1][0][1]
                # update columns
                new_column_tokens.append(new_sentence_tokens)
                new_column_concat.append(new_sentence[0][0])
                new_column_success_flag.append(1)
                new_column_pert_indices.append([token_index])
            # if token cannot be found in original list of tokens
            else:
                new_column_tokens.append(sentence)
                new_column_concat.append(df.iloc[s, orig_index])
                new_column_success_flag.append(0)
                new_column_pert_indices.append(None)

    df['change_first_name_concat'] = new_column_concat
    df['change_first_name_tokens'] = new_column_tokens
    df['change_first_name_success'] = new_column_success_flag
    df['change_first_name_pert_ind'] = new_column_pert_indices

def change_last_name(df: pd.DataFrame, orig_index: int, tokens_orig: list):
    """
    Change last name in sentence if one exists and if name is in CheckList's name lookup json

    :param df: DataFrame containing sentences
    :param orig_index: integer indicating the column index of the original sentence
    :param tokens_orig: tokenised version of sentence
    :return: None. Modifies DataFrame in-place
    """

    new_column_tokens, new_column_concat, new_column_success_flag, new_column_pert_indices = _gen_empty_columns()
    # Checklist requires pre-processing with Spacy for this perturbation
    pdata = list(nlp.pipe(df.iloc[:, orig_index]))

    for s in range(len(tokens_orig)):
        sentence = tokens_orig[s]
        new_sentence = Perturb.change_names(pdata[s], n=1, last_only=True, meta=True)
        if not new_sentence:
            new_column_tokens.append(sentence)
            new_column_concat.append(df.iloc[s, orig_index])
            new_column_success_flag.append(0)
            new_column_pert_indices.append(None)
        else:
            # extract token that has been perturbed
            token_pert = new_sentence[1][0][0]
            # verify that perturbed name appears as token in the original input
            if token_pert in sentence:
                # obtain index
                token_index = sentence.index(token_pert)
                new_sentence_tokens = deepcopy(sentence)
                new_sentence_tokens[token_index] = new_sentence[1][0][1]
                new_column_tokens.append(new_sentence_tokens)
                new_column_concat.append(new_sentence[0][0])
                new_column_success_flag.append(1)
                new_column_pert_indices.append([token_index])
            else:
                new_column_tokens.append(sentence)
                new_column_concat.append(df.iloc[s, orig_index])
                new_column_success_flag.append(0)
                new_column_pert_indices.append(None)

    df['change_last_name_concat'] = new_column_concat
    df['change_last_name_tokens'] = new_column_tokens
    df['change_last_name_success'] = new_column_success_flag
    df['change_last_name_pert_ind'] = new_column_pert_indices

def change_location(df: pd.DataFrame, orig_index: int, tokens_orig: list):
    """
    Change location in sentence if one exists and if location is in CheckList's location lookup json

    :param df: DataFrame containing sentences
    :param orig_index: integer indicating the column index of the original sentence
    :param tokens_orig: tokenised version of sentence
    :return: None. Modifies DataFrame in-place
    """

    new_column_tokens, new_column_concat, new_column_success_flag, new_column_pert_indices = _gen_empty_columns()
    # Checklist requires pre-processing with Spacy for this perturbation
    pdata = list(nlp.pipe(df.iloc[:, orig_index]))

    for s in range(len(tokens_orig)):
        sentence = tokens_orig[s]
        new_sentence = Perturb.change_location(pdata[s], n=1, meta=True)
        if not new_sentence:
            new_column_tokens.append(sentence)
            new_column_concat.append(df.iloc[s, orig_index])
            new_column_success_flag.append(0)
            new_column_pert_indices.append(None)
        else:
            # extract token that has been perturbed
            token_pert = new_sentence[1][0][0]
            # verify that perturbed name appears as token in the original input
            if token_pert in sentence:
                # obtain index
                token_index = sentence.index(token_pert)
                new_sentence_tokens = deepcopy(sentence)
                new_sentence_tokens[token_index] = new_sentence[1][0][1]
                new_column_tokens.append(new_sentence_tokens)
                new_column_concat.append(new_sentence[0][0])
                new_column_success_flag.append(1)
                new_column_pert_indices.append([token_index])
            # if token cannot be found in original list of tokens
            else:
                new_column_tokens.append(sentence)
                new_column_concat.append(df.iloc[s, orig_index])
                new_column_success_flag.append(0)
                new_column_pert_indices.append(None)

    df['change_location_concat'] = new_column_concat
    df['change_location_tokens'] = new_column_tokens
    df['change_location_success'] = new_column_success_flag
    df['change_location_pert_ind'] = new_column_pert_indices

def add_typo(df: pd.DataFrame, orig_index: int, tokens_orig: list):
    """
    Add typo using checklist library

    :param df: DataFrame containing sentences
    :param orig_index: integer indicating the column index of the original sentence
    :param tokens_orig: tokenised version of sentence
    :return: None. Modifies DataFrame in-place
    """

    new_column_tokens, new_column_concat, new_column_success_flag, new_column_pert_indices = _gen_empty_columns()

    for s in range(len(tokens_orig)):
        sentence = deepcopy(tokens_orig[s])
        if len(sentence) < 3:
            new_column_tokens.append(sentence)
            new_column_concat.append(df.iloc[s, orig_index])
            new_column_success_flag.append(0)
            new_column_pert_indices.append(None)
            continue
        token_length = 1
        attempts = 0
        while token_length <= 1 and attempts < 10:
            index_pert = random.randint(1, len(sentence) - 2)
            token_length = len(sentence[index_pert])
            attempts += 1
        if token_length <= 1:
            new_column_tokens.append(sentence)
            new_column_concat.append(df.iloc[s, orig_index])
            new_column_success_flag.append(0)
            new_column_pert_indices.append(None)
        else:
            sentence[index_pert + 1] = sentence[index_pert][-1] + sentence[index_pert + 1]
            sentence[index_pert] = sentence[index_pert][:-1]
            new_column_tokens.append(sentence)
            new_column_concat.append(" ".join(sentence))
            new_column_success_flag.append(1)
            new_column_pert_indices.append([index_pert, index_pert + 1])

    df['add_typo_concat'] = new_column_concat
    df['add_typo_tokens'] = new_column_tokens
    df['add_typo_success'] = new_column_success_flag
    df['add_typo_pert_ind'] = new_column_pert_indices

def strip_trailing_punct(df: pd.DataFrame, orig_index: int, tokens_orig: list):
    """
    Remove punctuation at end of sentence

    :param df: DataFrame containing sentences
    :param orig_index: integer indicating the column index of the original sentence
    :param tokens_orig: tokenised version of sentence
    :return: None. Modifies DataFrame in-place
    """

    new_column_tokens, new_column_concat, new_column_success_flag, new_column_pert_indices = _gen_empty_columns()

    for s in range(len(tokens_orig)):
        sentence = deepcopy(tokens_orig[s])
        if sentence[-1] in PUNCTUATION:
            sentence[-1] = ''
            new_column_concat.append(" ".join(sentence))
            new_column_success_flag.append(1)
            new_column_pert_indices.append([len(sentence)-1])
        else:
            new_column_concat.append(df.iloc[s, orig_index])
            new_column_success_flag.append(0)
            new_column_pert_indices.append(None)
        new_column_tokens.append(sentence)

    df['strip_punct_concat'] = new_column_concat
    df['strip_punct_tokens'] = new_column_tokens
    df['strip_punct_success'] = new_column_success_flag
    df['strip_punct_pert_ind'] = new_column_pert_indices

def remove_commas(df: pd.DataFrame, orig_index: int, tokens_orig: list):
    """
    Remove all commas from sentence

    :param df: DataFrame containing sentences
    :param orig_index: integer indicating the column index of the original sentence
    :param tokens_orig: tokenised version of sentence
    :return: None. Modifies DataFrame in-place
    """

    new_column_tokens, new_column_concat, new_column_success_flag, new_column_pert_indices = _gen_empty_columns()

    for s in range(len(tokens_orig)):
        sentence = tokens_orig[s]
        sentence_pert = _custom_remove_char(sentence, ',')
        if sentence == sentence_pert:
            new_column_tokens.append(tokens_orig[s])
            new_column_concat.append(df.iloc[s, orig_index])
            new_column_success_flag.append(0)
            new_column_pert_indices.append(None)
        else:
            empty_indices = []
            # tracking the indices of tokens that have been removed
            for t in range(len(sentence)):
                if sentence_pert[t] == '':
                    empty_indices.append(t)
            new_column_tokens.append(sentence_pert)
            new_column_concat.append(" ".join(sentence_pert))
            new_column_success_flag.append(1)
            new_column_pert_indices.append(empty_indices)

    df['remove_commas_concat'] = new_column_concat
    df['remove_commas_tokens'] = new_column_tokens
    df['remove_commas_success'] = new_column_success_flag
    df['remove_commas_pert_ind'] = new_column_pert_indices

def remove_all_punctuation(df: pd.DataFrame, orig_index: int, tokens_orig: list):
    """
    Remove all punctuation from sentence

    :param df: DataFrame containing sentences
    :param orig_index: integer indicating the column index of the original sentence
    :param tokens_orig: tokenised version of sentence
    :return: None. Modifies DataFrame in-place
    """

    new_column_tokens, new_column_concat, new_column_success_flag, new_column_pert_indices = _gen_empty_columns()

    for s in range(len(tokens_orig)):
        sentence = tokens_orig[s]
        sentence_pert = _custom_remove_char(sentence, PUNCTUATION)
        if sentence == sentence_pert:
            new_column_tokens.append(tokens_orig[s])
            new_column_concat.append(df.iloc[s, orig_index])
            new_column_success_flag.append(0)
            new_column_pert_indices.append(None)
        else:
            empty_indices = []
            # tracking the indices of tokens that have been removed
            for t in range(len(sentence)):
                if sentence_pert[t] == '':
                    empty_indices.append(t)
            new_column_tokens.append(sentence_pert)
            new_column_concat.append(" ".join(sentence_pert))
            new_column_success_flag.append(1)
            new_column_pert_indices.append(empty_indices)

    df['remove_all_punct_concat'] = new_column_concat
    df['remove_all_punct_tokens'] = new_column_tokens
    df['remove_all_punct_success'] = new_column_success_flag
    df['remove_all_punct_pert_ind'] = new_column_pert_indices

def switch_gender(df: pd.DataFrame, orig_index: int, tokens_orig: list):
    """
    Change gendered words

    :param df: DataFrame containing sentences
    :param orig_index: integer indicating the column index of the original sentence
    :param tokens_orig: tokenised version of sentence
    :return: None. Modifies DataFrame in-place
    """

    new_column_tokens, new_column_concat, new_column_success_flag, new_column_pert_indices = _gen_empty_columns()

    for s in range(len(tokens_orig)):
        sentence = deepcopy(tokens_orig[s])
        perturbed_indices = []
        for t in range(len(sentence)):
            if sentence[t] in DICT_GENDER:
                sentence[t] = DICT_GENDER[sentence[t]]
                perturbed_indices.append(t)
        new_column_tokens.append(sentence)
        if len(perturbed_indices) == 0:
            new_column_concat.append(df.iloc[s, orig_index])
            new_column_success_flag.append(0)
            new_column_pert_indices.append(None)
        else:
            new_column_concat.append(" ".join(sentence))
            new_column_success_flag.append(1)
            new_column_pert_indices.append(perturbed_indices)

    df['switch_gender_concat'] = new_column_concat
    df['switch_gender_tokens'] = new_column_tokens
    df['switch_gender_success'] = new_column_success_flag
    df['switch_gender_pert_ind'] = new_column_pert_indices

def add_perturbations(
        df: pd.DataFrame, tokenizer, sentence_col_name: str, perturbation_functions, seed=3
):
    """
    Apply multiple perturbations, generating a new column for each perturbation

    :param tokenizer: SpacyTokenizer or BertTokenizer
    :param df: DataFrame containing sentences
    :param sentence_col_name: Name of column containing sentence to be perturbed
    :param perturbation_functions: List of perturbation functions
    :param seed: Random seed
    :return: DataFrame with additional columns containing perturbed sentences and success flags
    """
    df = df.copy()

    tokenizer = tokenizer or SpacyTokenizer()

    orig_index = df.columns.get_loc(sentence_col_name)

    tokens_orig = [[str(x) for x in tokenizer.tokenize(df.iloc[i, orig_index])] for i in range(len(df))]

    df['tokens_orig'] = tokens_orig

    np.random.seed(seed)  # Set seed as some perturbations are stochastic

    for perturbation in perturbation_functions:
        perturbation(df, orig_index, tokens_orig)

    return df