text:
import logging
import pathlib
import json
import time
import torch
from torch import utils
import numpy as np
from train import select_device, TRAIN_DTYPES
from train import create_dataset as train_create_dataset
import methods
import dataset
import integrators
from systems import spring, wave, spring_mesh, navier_stokes
from collections import namedtuple
import math

class TrajTimeCollector:
    def __init__(self):
        self.times = {}
        self.current_traj = None

    def start_traj(self, traj_num):
        self.current_traj = traj_num
        if self.current_traj not in self.times:
            self.times[self.current_traj] = 0

    def accumulate_time(self, time):
        self.times[self.current_traj] += time

    def get_time(self, traj_num):
        return self.times.get(traj_num)

def load_network(net_dir, base_dir, base_logger, model_file_name="model.pt"):
    logger = base_logger.getChild("load_network")
    if net_dir is None:
        logger.info("Skipping network load")
        return None
    # Load metadata
    net_dir = base_dir / pathlib.Path(net_dir)
    meta_path = net_dir / "model.json"
    with open(meta_path, "r", encoding="utf8") as meta_file:
        metadata = json.load(meta_file)
    logger.info(f"Loaded network description from {meta_path}")
    # Create network
    net = methods.build_network(metadata)
    # Load weights
    weight_path = net_dir / model_file_name
    net.load_state_dict(torch.load(weight_path, map_location="cpu"))
    logger.info(f"Loaded weights from {weight_path}")
    return net

def create_dataset(base_dir, data_args):
    dataset_type = data_args.get("dataset", None)
    if dataset_type is None:
        data_dir = base_dir / data_args["data_dir"]
        linearize = data_args.get("linearize", False)
        data_set = dataset.TrajectoryDataset(data_dir=data_dir, linearize=linearize)
        loader = utils.data.DataLoader(data_set, batch_size=1, shuffle=False)
        return data_set, loader
    else:
        # Non-default data set for evaluation
        # Use this for creating more complex datasets, like pytorch-geometric
        train_dataset, train_loader, _val_dataset, _val_loader = train_create_dataset(base_dir=base_dir, data_args=data_args)
        return train_dataset, train_loader

def raw_err(approx, true, norm=2):
    err = np.linalg.norm(approx - true, ord=norm, axis=1)
    return err

def rel_err(approx, true, norm=2):
    num = np.linalg.norm(approx - true, ord=norm, axis=1)
    denom = np.linalg.norm(true, ord=norm, axis=1)
    return num / denom

def mean_square_err(approx, true):
    diff = (approx - true)**2
    return np.mean(diff, axis=1)

def train_knn(net, eval_args, base_dir, base_logger):
    # Load the training data
    logger = base_logger.getChild("knn_train")
    eval_type = eval_args["eval_type"]
    if eval_type == "knn-predictor-oneshot":
        assert eval_args["train_data"]["dataset"] == "step-snapshot"
        assert eval_args["train_data"]["loader"].get("type", "pytorch") == "pytorch"
        assert eval_args["train_data"]["loader"]["batch_size"] == 1

    train_dataset, train_loader = create_dataset(base_dir, eval_args["train_data"])
    # Train the KNN
    if eval_type == "knn-regressor-oneshot":
        logger.info("Starting fitting of dataset for KNN Regressor.")
        data_x = []
        data_y = []
        for batch in train_loader:
            data_x.append(np.concatenate([batch.p, batch.q], axis=-1))
            data_y.append(np.concatenate([batch.dp_dt, batch.dq_dt], axis=-1))
        data_x = np.concatenate(data_x, axis=0)
        data_y = np.concatenate(data_y, axis=0)
        net.fit(data_x, data_y)
        logger.info("Finished fitting of dataset for KNN Regressor.")
    elif eval_type == "knn-predictor-oneshot":
        logger.info("Starting fitting of dataset for KNN Predictor.")
        data_x = []
        data_y = []
        for batch in train_loader:
            assert batch.p.shape[0] == 1
            assert batch.q.shape[0] == 1
            assert batch.p_noiseless.shape[0] == 1
            assert batch.q_noiseless.shape[0] == 1
            data_x.append(np.concatenate([batch.p[0], batch.q[0]], axis=-1))
            data_y.append(np.concatenate([batch.p_step[0], batch.q_step[0]], axis=-1))
        data_x = np.stack(data_x, axis=0)
        data_y = np.stack(data_y, axis=0)
        net.fit(data_x, data_y)
        logger.info("Finished fitting of dataset for KNN Predictor.")
    # Return the KNN
    return net

def run_phase(base_dir, out_dir, phase_args):
    logger = logging.getLogger("eval")
    base_dir = pathlib.Path(base_dir)
    out_dir = pathlib.Path(out_dir)
    eval_args = phase_args["eval"]
    eval_only_time_collector = TrajTimeCollector()

    # Choose the device and dtype
    device = select_device(try_gpu=eval_args["try_gpu"], base_logger=logger)
    eval_type = eval_args["eval_type"]
    eval_dtype, eval_dtype_np = TRAIN_DTYPES[eval_args.get("eval_dtype", "float")]
    baseline_stride = int(eval_args.get("coarsening", 1))

    # Load the network
    if eval_type in {"knn-regressor-oneshot", "knn-predictor-oneshot"}:
        logger.info(f"Doing \"oneshot\" KNN training: {eval_type}")
        net = methods.build_network(net_args={"arch": eval_type, "arch_args": None})
    else:
        # Normal network load process
        net = load_network(phase_args["eval_net"], base_dir=base_dir, base_logger=logger,
                           model_file_name=phase_args.get("eval_net_file", "model.pt"))

    # Send network to correct device if PyTorch
    if isinstance(net, (torch.nn.Module, torch.Tensor)):
        net = net.to(device, dtype=eval_dtype)
        net.eval()

    # Handle possible knn "oneshot" training before evaluation
    if eval_type in {"knn-regressor-oneshot", "knn-predictor-oneshot"}:
        net = train_knn(net=net, eval_args=eval_args,
                        base_dir=base_dir, base_logger=logger)

    # Load the evaluation data
    eval_dataset, eval_loader = create_dataset(base_dir, phase_args["eval_data"])

    # Integrate each trajectory, compute stats and store
    logger.info("Starting evaluation")
    integrator_type = eval_args["integrator"]

    if eval_type in {"mlp-deriv", "nn-kernel-deriv"}:
        def mlp_time_deriv_func(batch):
            MLPDerivative = namedtuple("MLPDerivative", ["dq_dt", "dp_dt"])

            extra_data = batch.extra_fixed_mask
            if torch.is_tensor(extra_data):
                extra_data = extra_data[0].to(device, dtype=eval_dtype)
                extra_data.requires_grad = False
            else:
                extra_data = None

            def net_no_grad(q, p, dt=1.0, t=0):
                with torch.no_grad():
                    q = torch.from_numpy(q).to(device, dtype=eval_dtype)
                    p = torch.from_numpy(p).to(device, dtype=eval_dtype)
                    if device.type == "cuda":
                        torch.cuda.synchronize()
                    t_start = time.perf_counter()
                    ret = net(q=q, p=p, extra_data=extra_data)
                    if device.type == "cuda":
                        torch.cuda.synchronize()
                    t_end = time.perf_counter()
                    eval_only_time_collector.accumulate_time(t_end - t_start)
                    dq = ret.dq_dt.detach().cpu().numpy()
                    dp = ret.dp_dt.detach().cpu().numpy()
                    return MLPDerivative(dq_dt=dq, dp_dt=dp)

            return net_no_grad

    elif eval_type in {"mlp-step", "nn-kernel-step"}:
        def mlp_time_deriv_func(batch):
            MLPPrediction = namedtuple("MLPPrediction", ["q", "p"])

            extra_data = batch.extra_fixed_mask
            if torch.is_tensor(extra_data):
                extra_data = extra_data[0].to(device, dtype=eval_dtype)
                extra_data.requires_grad = False
            else:
                extra_data = None

            def model_next_step(q, p, dt=1.0, t=0):
                with torch.no_grad():
                    q = torch.from_numpy(q).to(device, dtype=eval_dtype)
                    p = torch.from_numpy(p).to(device, dtype=eval_dtype)
                    if device.type == "cuda":
                        torch.cuda.synchronize()
                    t_start = time.perf_counter()
                    ret = net(q=q, p=p, extra_data=extra_data)
                    if device.type == "cuda":
                        torch.cuda.synchronize()
                    t_end = time.perf_counter()
                    eval_only_time_collector.accumulate_time(t_end - t_start)
                    next_q = ret.q.detach().cpu().numpy()
                    next_p = ret.p.detach().cpu().numpy()
                    return MLPPrediction(q=next_q, p=next_p)

            return model_next_step
        if integrator_type != "null":
            raise ValueError(f"mlp-step predictions do not work with integrator {integrator_type}")
    elif eval_type in {"cnn-deriv", "unet-deriv"}:
        def cnn_time_deriv_func(batch):
            CNNDerivative = namedtuple("CNNDerivative", ["dq_dt", "dp_dt"])

            p_full_shape = (1, ) + batch.p[0, 0].shape
            q_full_shape = (1, ) + batch.q[0, 0].shape

            if len(q_full_shape) < 3:
                q_full_shape = q_full_shape + (1, )
            if len(p_full_shape) < 3:
                p_full_shape = p_full_shape + (1, )

            extra_data = batch.extra_fixed_mask
            if torch.is_tensor(extra_data):
                extra_data = extra_data[0].to(device, dtype=eval_dtype)
                if extra_data.ndim < 3:
                    extra_data = extra_data.unsqueeze(-1)
                extra_data.requires_grad = False
            else:
                extra_data = None

            def cnn_time_deriv(q, p, dt=1.0, t=0):
                with torch.no_grad():
                    q_orig_shape = q.shape
                    p_orig_shape = p.shape
                    q = torch.from_numpy(q).to(device, dtype=eval_dtype).reshape(q_full_shape)
                    p = torch.from_numpy(p).to(device, dtype=eval_dtype).reshape(p_full_shape)
                    if device.type == "cuda":
                        torch.cuda.synchronize()
                    t_start = time.perf_counter()
                    res = net(p=p, q=q, extra_data=extra_data)
                    if device.type == "cuda":
                        torch.cuda.synchronize()
                    t_end = time.perf_counter()
                    eval_only_time_collector.accumulate_time(t_end - t_start)
                    dq_dt = res.dq_dt.reshape(q_orig_shape).cpu().numpy()
                    dp_dt = res.dp_dt.reshape(p_orig_shape).cpu().numpy()
                    return CNNDerivative(dq_dt=dq_dt, dp_dt=dp_dt)

            return cnn_time_deriv

    elif eval_type in {"cnn-step", "unet-step"}:
        def cnn_time_deriv_func(batch):
            CNNPrediction = namedtuple("CNNPrediction", ["q", "p"])

            p_full_shape = (1, ) + batch.p[0, 0].shape
            q_full_shape = (1, ) + batch.q[0, 0].shape

            if len(q_full_shape) < 3:
                q_full_shape = q_full_shape + (1, )
            if len(p_full_shape) < 3:
                p_full_shape = p_full_shape + (1, )

            extra_data = batch.extra_fixed_mask
            if torch.is_tensor(extra_data):
                extra_data = extra_data[0].to(device, dtype=eval_dtype)
                if extra_data.ndim < 3:
                    extra_data = extra_data.unsqueeze(-1)
                extra_data.requires_grad = False
            else:
                extra_data = None

            def model_next_step(q, p, dt=1.0, t=0):
                with torch.no_grad():
                    q_orig_shape = q.shape
                    p_orig_shape = p.shape
                    q = torch.from_numpy(q).to(device, dtype=eval_dtype).reshape(q_full_shape)
                    p = torch.from_numpy(p).to(device, dtype=eval_dtype).reshape(p_full_shape)
                    if device.type == "cuda":
                        torch.cuda.synchronize()
                    t_start = time.perf_counter()
                    res = net(p=p, q=q, extra_data=extra_data)
                    if device.type == "cuda":
                        torch.cuda.synchronize()
                    t_end = time.perf_counter()
                    eval_only_time_collector.accumulate_time(t_end - t_start)
                    q_next = res.q.reshape(q_orig_shape).cpu().numpy()
                    p_next = res.p.reshape(p_orig_shape).cpu().numpy()
                    return CNNPrediction(q=q_next, p=p_next)

            return model_next_step

        if integrator_type != "null":
            raise ValueError(f"cnn-step predictions do not work with integrator {integrator_type}")

    elif eval_type in {"knn-regressor", "knn-regressor-oneshot"}:
        # Use the time_derivative
        KNNDerivative = namedtuple("KNNDerivative", ["dq_dt", "dp_dt"])
        def model_time_deriv(q, p, dt=1.0, t=0):
            p_len = p.shape[-1]
            x = np.concatenate([p, q], axis=-1)
            t_start = time.perf_counter()
            ret = net.predict(x)
            t_end = time.perf_counter()
            eval_only_time_collector.accumulate_time(t_end - t_start)
            dpdt, dqdt = np.split(ret, [p_len], axis=-1)
            return KNNDerivative(dq_dt=dqdt, dp_dt=dpdt)
        time_deriv_func = model_time_deriv
    elif eval_type in {"knn-predictor", "knn-predictor-oneshot"}:
        KNNPrediction = namedtuple("KNNPrediction", ["q", "p"])
        def model_next_step(q, p, dt=1.0, t=0):
            p_len = p.shape[-1]
            x = np.concatenate([p, q], axis=-1)
            t_start = time.perf_counter()
            ret = net.predict(x)
            t_end = time.perf_counter()
            eval_only_time_collector.accumulate_time(t_end - t_start)
            next_p, next_q = np.split(ret, [p_len], axis=-1)
            return KNNPrediction(q=next_q, p=next_p)
        time_deriv_func = model_next_step
        if integrator_type != "null":
            raise ValueError(f"KNN predictions do not work with integrator {integrator_type}")
    elif eval_type == "integrator-baseline":
        # Use the time_derivative
        SystemDerivative = namedtuple("SystemDerivative", ["dq_dt", "dp_dt"])
        def system_derivative(q, p, dt=1.0, t=0):
            if torch.is_tensor(dt):
                dt = dt.item()
            t_start = time.perf_counter()
            dq_dt, dp_dt = system.derivative(p=p, q=q)
            t_end = time.perf_counter()
            eval_only_time_collector.accumulate_time(t_end - t_start)
            return SystemDerivative(dp_dt=dp_dt, dq_dt=dq_dt)
        time_deriv_func = system_derivative
    else:
        logger.error(f"Invalid evaluation type: {eval_type}")
        raise ValueError(f"Invalid evaluation type: {eval_type}")

    trajectory_results = []

    for traj_num, trajectory in enumerate(eval_loader):
        logger.info(f"Starting trajectory number {traj_num}")
        eval_only_time_collector.start_traj(traj_num)
        traj_name = trajectory.name[0]
        p = trajectory.p.to(device, dtype=eval_dtype)
        q = trajectory.q.to(device, dtype=eval_dtype)
        p_noiseless = trajectory.p_noiseless.to(device, dtype=eval_dtype)

        q_noiseless = trajectory.q_noiseless.to(device, dtype=eval_dtype)
        masses = trajectory.masses.to(device, dtype=eval_dtype)
        num_time_steps = int(trajectory.trajectory_meta["num_time_steps"][0].detach().cpu().item())
        time_step_size = float(trajectory.trajectory_meta["time_step_size"][0].detach().cpu().item())
        edges = None
        vertices = None
        boundary_cond_func = None
        static_nodes = None

        # Compute hamiltonians
        # Construct systems
        if eval_dataset.system == "spring":
            system = spring.system_from_records()
        elif eval_dataset.system == "wave":
            n_grid = eval_dataset.system_metadata["n_grid"]
            space_max = eval_dataset.system_metadata["space_max"]
            wave_speed = trajectory.trajectory_meta["wave_speed"][0].item()
            system = wave.system_from_records(n_grid=n_grid,
                                              space_max=space_max,
                                              wave_speed=wave_speed)
        elif eval_dataset.system == "spring-mesh":
            n_dim = eval_dataset.system_metadata["n_dim"]
            particles = eval_dataset.system_metadata["particles"]
            n_particles = len(eval_dataset.system_metadata["particles"])
            edges_dict = eval_dataset.system_metadata["edges"]
            edges = np.array([(e["a"], e["b"]) for e in edges_dict] +
                             [(e["b"], e["a"]) for e in edges_dict], dtype=np.int64).T
            vel_decay = eval_dataset.system_metadata["vel_decay"]
            system = spring_mesh.system_from_records(n_dim, particles, edges_dict,
                                                     vel_decay=vel_decay)
            boundary_cond_func = spring_mesh.make_enforce_boundary_function(trajectory=trajectory)
        elif eval_dataset.system == "navier-stokes":
            grid_resolution = eval_dataset.system_metadata["grid_resolution"]
            in_velocity = trajectory.trajectory_meta["in_velocity"][0].item()
            viscosity = trajectory.trajectory_meta["viscosity"][0].item()
            vertices = eval_dataset[0].vertices.tolist()
            vertices_array = eval_dataset[0].vertices
            vertices_array.setflags(write=False)
            edges = eval_dataset[0].edge_index.tolist()
            system = navier_stokes.system_from_records(grid_resolution=grid_resolution, viscosity=viscosity)

            fixed_mask_solutions = trajectory.fixed_mask_p[0].cpu().numpy()
            fixed_mask_solutions.setflags(write=False)
            fixed_mask_pressures = trajectory.fixed_mask_q[0].cpu().numpy()
            fixed_mask_pressures.setflags(write=False)
            boundary_cond_func = navier_stokes.make_enforce_boundary_function(
                in_velocity=in_velocity,
                vertex_coords=vertices_array,
                fixed_mask_solutions=fixed_mask_solutions,
                fixed_mask_pressures=fixed_mask_pressures,
            )
            static_nodes = trajectory.static_nodes.numpy().squeeze()
        else:
            raise ValueError(f"Unknown system type {eval_dataset.system}")

        if eval_type in {"cnn-step", "cnn-deriv", "unet-step", "unet-deriv"}:
            time_deriv_func = cnn_time_deriv_func(batch=trajectory)
        elif eval_type in {"mlp", "mlp-step", "mlp-deriv", "nn-kernel-step", "nn-kernel-deriv"}:
            time_deriv_func = mlp_time_deriv_func(batch=trajectory)

        p0 = p[:, 0].detach().cpu().numpy()
        q0 = q[:, 0].detach().cpu().numpy()
        t0 = trajectory.t[0, 0].item()
        q0, p0 = q0, p0
        integrate_start = time.perf_counter()
        if not getattr(eval_dataset, "_linearize", True):
            q0 = q0.reshape((1, -1))
            p0 = p0.reshape((1, -1))

        if eval_type == "integrator-baseline" and baseline_stride != 1:
            num_time_steps = math.ceil(num_time_steps / baseline_stride)
            time_step_size = time_step_size * baseline_stride
            logger.info(f"Trajectory coarsened to {time_step_size} steps")

        int_res_raw = integrators.numerically_integrate(
            integrator=integrator_type,
            q0=q0,
            p0=p0,
            t0=t0,
            num_steps=num_time_steps,
            dt=time_step_size,
            deriv_func=time_deriv_func,
            boundary_cond_func=boundary_cond_func,
            system=system)

        # Remove extraneous batch dimension
        integrate_elapsed = time.perf_counter() - integrate_start
        # Split the integration result
        int_p = int_res_raw.p
        int_q = int_res_raw.q

        if not getattr(eval_dataset, "_linearize", True):
            n_noiseless_steps = p_noiseless.shape[1]
            p_noiseless = p_noiseless.reshape((1, n_noiseless_steps, -1))
            q_noiseless = q_noiseless.reshape((1, n_noiseless_steps, -1))

        # Compute errors and other statistics
        int_res = np.concatenate([int_p, int_q], axis=-1)
        true = torch.cat([p_noiseless, q_noiseless], axis=-1)[0].detach().cpu().numpy()

        if eval_type == "integrator-baseline" and baseline_stride != 1:
            true = true[::baseline_stride]

        raw_l2 = raw_err(approx=int_res, true=true, norm=2)
        rel_l2 = rel_err(approx=int_res, true=true, norm=2)
        mse_err = mean_square_err(approx=int_res, true=true)

        mean_raw_l2 = np.mean(raw_l2)
        mean_mse = np.mean(mse_err)
        logger.info(f"Mean raw l2: {mean_raw_l2}")
        logger.info(f"Mean MSE: {mean_mse}")

        int_p = int_res_raw.p
        int_q = int_res_raw.q

        # All combinations true/net hamiltonian on true/net trajectories

        int_numpy_arrays = {
            traj_name: int_res,
            f"{traj_name}_relerr_l2": rel_l2,
            f"{traj_name}_raw_l2": raw_l2,
            f"{traj_name}_mse": mse_err,
            f"{traj_name}_p": int_p,
            f"{traj_name}_q": int_q,
        }

        int_stats = {
            "name": traj_name,
            "file_names": {
                "traj": traj_name,
                "relerr_l2": f"{traj_name}_relerr_l2",
                "raw_l2": f"{traj_name}_raw_l2",
                "mse": f"{traj_name}_mse",
                "p": f"{traj_name}_p",
                "q": f"{traj_name}_q",
            },
            "timing": {
                "integrate_elapsed": integrate_elapsed,
                "eval_only": eval_only_time_collector.get_time(traj_num),
            }
        }

        trajectory_results.append((traj_name, int_numpy_arrays, int_stats))
        logger.info(f"Integrated trajectory: {traj_num} ({traj_name}) in {integrate_elapsed} sec")

    # Process for save
    results_meta = {"integration_stats": []}
    saved_trajectories = {}
    for _traj_name, int_res, int_stats in trajectory_results:
        saved_trajectories.update(int_res)
        results_meta["integration_stats"].append(int_stats)

    # Save integration results
    int_traj_file = out_dir / "integrated_trajectories.npz"
    np.savez(int_traj_file, **saved_trajectories)
    logger.info(f"Saved integrated trajectories to {int_traj_file}")

    # Save metadata
    int_meta_file = out_dir / "results_meta.json"
    with open(int_meta_file, "w", encoding="utf8") as meta_file:
        json.dump(results_meta, meta_file)
    logger.info(f"Saved trajectory integration metadata to {int_meta_file}")

max_stars_repo_name: karlotness/nn-benchmark

max_issues_repo_name: karlotness/nn-benchmark

max_forks_repo_name: karlotness/nn-benchmark


text:
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import  average_precision_score, roc_auc_score

from optim.loss import EarlyStopping
import os

def train(args, logger, data, model, path):
    checkpoints_path = path

    logger.info('Start training')
    logger.info(
        f'dropout:{args.dropout}, nu:{args.nu},seed:{args.seed},lr:{args.lr},self-loop:{args.self_loop},norm:{args.norm}')

    logger.info(
        f'n-epochs:{args.n_epochs}, n-hidden:{args.n_hidden},n-layers:{args.n_layers},weight-decay:{args.weight_decay}')

    optimizer = torch.optim.Adam(model.parameters(),
                                 lr=args.lr,
                                 weight_decay=args.weight_decay)
    if args.early_stop:
        stopper = EarlyStopping(patience=100)

    adj = data['g'].adjacency_matrix().to_dense().cuda()
    loss_fn = nn.MSELoss()
    model.train()

    # 创立矩阵以存储结果曲线
    arr_epoch = np.arange(args.n_epochs)
    arr_loss = np.zeros(args.n_epochs)
    arr_valauc = np.zeros(args.n_epochs)
    arr_testauc = np.zeros(args.n_epochs)
    savedir = './embeddings/Dominant/' + args.dataset + '/{}'.format(args.abclass)
    if not os.path.exists(savedir):
        os.makedirs(savedir)
    max_valauc = 0
    epoch_max = 0
    dur = []
    for epoch in range(args.n_epochs):
        # model.train()
        # if epoch %5 == 0:
        t0 = time.time()
        # forward
        z, re_x, re_adj = model(data['g'], data['features'])
        if args.module == 'Dominant':
            loss = Recon_loss(re_x, re_adj, adj, data['features'], data['train_mask'], loss_fn, 'AX', 0.2)
        else:
            loss = Recon_loss(re_x, re_adj, adj, data['features'], data['train_mask'], loss_fn, 'A', 0)

        # 保存训练loss
        arr_loss[epoch] = loss.item()
        #

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if epoch >= 3:
            dur = time.time() - t0

        auc, ap, val_loss = fixed_graph_evaluate(args, model, data, adj, data['val_mask'],0.2)
        # 保存验证集AUC
        arr_valauc[epoch] = auc
        # 保存验证集AUC

        if auc > max_valauc:
            max_valauc = auc
            epoch_max = epoch
            torch.save(model.state_dict(), checkpoints_path)
            np.save(os.path.sep.join(
                [savedir + '/embeddings_max.npy']),
                z[data['test_mask']].data.cpu().numpy())
            np.save(os.path.sep.join([savedir + '/label_max.npy'.format(args.n_epochs)]),
                    data['labels'][data['test_mask']].cpu().numpy())

        print(
            "Epoch {:05d} | Time(s) {:.4f} | Train Loss {:.4f}  | Val AUROC {:.4f}  | "
            "ETputs(KTEPS) {:.2f}".format(epoch, np.mean(dur), loss.item() * 100000,
                                          auc,  data['n_edges'] / np.mean(dur) / 1000))

    if args.early_stop:
        print('loading model before testing.')
        model.load_state_dict(torch.load(checkpoints_path))

        # if epoch%100 == 0:

    auc, ap, _ = fixed_graph_evaluate(args, model, data, adj, data['test_mask'],0.2)
    test_dur = 0
    # 保存测试集AUC
    arr_testauc[epoch] = auc
    # 保存测试集AUC
    best_epoch = epoch_max
    print("Test Time {:.4f} | Test AUROC {:.4f} | Test AUPRC {:.4f}".format(test_dur, auc, ap))

    return auc, ap, best_epoch

def Recon_loss(re_x, re_adj, adj, x, mask, loss_fn, mode,alpha):
    # S_loss: structure loss A_loss: Attribute loss
    if mode == 'A':
        return loss_fn(re_x[mask], x[mask])
    if mode == 'X':
        return loss_fn(re_adj[mask], adj[mask])
    if mode == 'AX':
        return (1-alpha)*loss_fn(re_x[mask], x[mask]) + alpha*loss_fn(re_adj[mask], adj[mask])

def anomaly_score(re_x, re_adj, adj, x, mask, loss_fn, mode,alpha):
    if mode == 'A':
        A_scores = F.mse_loss(re_x[mask], x[mask], reduction='none')
        return torch.mean(A_scores, 1)
    if mode == 'X':
        S_scores = F.mse_loss(re_adj[mask], adj[mask], reduction='none')
        return torch.mean(S_scores, 1)

    if mode == 'AX':
        A_scores = F.mse_loss(re_x[mask], x[mask], reduction='none')
        S_scores = F.mse_loss(re_adj[mask], adj[mask], reduction='none')
        return (1-alpha)*torch.mean(A_scores, 1) + alpha*torch.mean(S_scores, 1)

def fixed_graph_evaluate(args, model, data, adj, mask,alpha):
    loss_fn = nn.MSELoss()

    model.eval()
    with torch.no_grad():
        labels = data['labels'][mask]

        loss_mask = mask.bool() & data['labels'].bool()
        z, re_x, re_adj = model(data['g'], data['features'])
        if args.module == 'GAE':
            loss = Recon_loss(re_x, re_adj, adj, data['features'], loss_mask, loss_fn, 'X', alpha)
            scores = anomaly_score(re_x, re_adj, adj, data['features'], mask, loss_fn, 'X', alpha)
        else:
            loss = Recon_loss(re_x, re_adj, adj, data['features'], loss_mask, loss_fn, 'AX',alpha)

            scores = anomaly_score(re_x, re_adj, adj, data['features'], mask, loss_fn, 'AX',alpha)

        labels = labels.cpu().numpy()
        scores = scores.cpu().numpy()
        auc = roc_auc_score(labels, scores)
        ap = average_precision_score(labels, scores)

    return auc, ap, loss

ext: py

max_stars_repo_name: haoyfan/Dual-SVDAE

max_issues_repo_name: haoyfan/Dual-SVDAE

max_forks_repo_name: haoyfan/Dual-SVDAE

reason: import numpy


text:
import gc
import logging
import numpy as np
import datetime
from queue import Queue
from multiprocessing import Process
from collections import deque

from model import *
from hyper_params import *
from game import *
import go_board

def evaluate(model_a, model_b, num_games=1):
    logging.info("Evaluation start:")
    results = []
    for eval_game_index in np.arange(num_games):
        start_player, outcome = two_player_play(eval_game_index, model_a, model_b)
        results.append((start_player, outcome))
    logging.info("Evaluation end: results: %s", results)
    return results

def train_loop():
    best_model = SimpleNNModel()
    trained_model = best_model.copy()

    np.random.seed(SEED)

    history = deque(maxlen=MAX_GAMES_HISTORY_SIZE)
    #game_history is a list of tuples (pi, s, z)
    # pi - the probability distribution over actions at state s
    # s - state
    # z - outcome of the game

    for iter_index in np.arange(NUM_ITER):
        logging.debug("Iter: %d", iter_index)

        for game_index in range(NUM_GAMES):
            self_play(game_index, best_model, history, noise_alpha=DIRICHLET_ALPHA, temperatures=TEMPERATURES)

        for train_iter_idx in np.arange(NUM_TRAIN_LOOP_ITER):
            logging.debug("Network train: %d", train_iter_idx)
            batch = np.random.choice(history, size=BATCH_SIZE)
            trained_model.train_on_hist_batch(batch, train_iter_idx)

        eval_results = evaluate(trained_model, best_model, num_games=NUM_EVAL_GAMES)
        outcomes = np.array([outcome for _, outcome in eval_results])
        trained_model_win_ratio = np.sum(outcomes == OUTCOME_WIN_PLAYER_1)/len(outcomes)

        logging.info("Iter %d evaluation: trained_model_win percent : %.2f%%", iter_index, trained_model_win_ratio*100)

        if trained_model_win_ratio > MIN_MODEL_REPLACEMENT_WIN_RATIO:
            old_best = best_model
            best_model = trained_model.copy()
            best_model.save("%s/model-%s-%d-%.0f.h5" % (OUTPUT_DIR, log_id, iter_index, trained_model_win_ratio*100))
            del old_best

        gc.collect()

if __name__ == "__main__":
    log_id = "ago-%dx%dx%d-%s" % (BOARD_SIZE, BOARD_SIZE, MCTS_STEPS, datetime.datetime.now().strftime("%m%d%H%M"))
    log_file_handler = logging.FileHandler('%s/%s.log' % (OUTPUT_DIR, log_id))

    logging.basicConfig(level=logging.DEBUG)
    logger = logging.getLogger()
    logger.addHandler(log_file_handler)

    train_loop()

max_stars_repo_name: rocipher/alphago-zero

max_issues_repo_name: rocipher/alphago-zero

max_forks_repo_name: rocipher/alphago-zero


text:
import copy
import numpy as np
import matplotlib.pyplot as plt

class BaseSprite:
    def __init__(self, position=(0, 0)):
        self.position = np.asarray(position)  # row, col

    def __setattr__(self, name, value):
        if name == 'position':
            value = copy.deepcopy(np.asarray(value, dtype=int))
        super().__setattr__(name, value)

    def plot(self, ax):
        xy = pos2xy(self.position) + (0.5, 0.5)
        c = plt.Circle(xy, 0.2, color='k', fill=False, linewidth=1)
        ax.add_patch(c)

def pos2xy(pos):
    return np.asarray(pos)[::-1]

max_stars_repo_name: abagaria/visgrid

max_issues_repo_name: abagaria/visgrid

max_forks_repo_name: abagaria/visgrid


text:
import numpy as __np__
from numpy import cos as __cos__
from numpy import sin as __sin__
from numpy import arctan2 as __arctan2__
from numpy import sqrt as __sqrt__
import matplotlib.pyplot as __plt__
from matplotlib import cm as __cm__
from . import tools as __tools__

class Coefficient(object):
	"""
	Return a set of Seidel wavsfront aberrations Coefficient
	"""
	__coefficients__ = []
	__seidellist___=["Ap Piston",
					"At Tilt",
					"Ad Defocus",
		 			"As Astigmatism",
		 			"Ac Coma",
		 			"As Spherical"]
	def __init__(self,Ap=0,Bp=0,At=0,Bt=0,Ad=0,Bd=0,Aa=0,Ba=0,Ac=0,Bc=0,As=0,Bs=0):
		if type(Ap) == list:
			self.__coefficients__ = Ap
		else:
			self.__coefficients__ = [[Ap,Bp],[At,Bt],[Ad,Bd],[Aa,Ba],[Ac,Bc],[As,Bs]]
	def outputcoefficient(self):
		return self.__coefficients__

	def seidelsurface(self, label = True, zlim=[], matrix = False):
		r1 = __np__.linspace(0, 1, 100)
		u1 = __np__.linspace(0, 2*__np__.pi, 100)
		[u,r] = __np__.meshgrid(u1,r1)
		X = r*__cos__(u)
		Y = r*__sin__(u)
		W = __seidelpolar__(self.__coefficients__,r,u)
		fig = __plt__.figure(figsize=(12, 8), dpi=80)
		ax = fig.gca(projection='3d')
		surf = ax.plot_surface(X, Y, W, rstride=1, cstride=1, cmap=__cm__.RdYlGn,
	        linewidth=0, antialiased=False, alpha = 0.6)
		fig.colorbar(surf, shrink=1, aspect=30)
		__plt__.show()

	def twyman_green(self, lambda_1 = 632, PR = 1):
		lambda_1 = lambda_1*(10**-9)
		A = self.__coefficients__
		r = __np__.linspace(-PR, PR, 400)
		x, y = __np__.meshgrid(r,r)
		OPD = __seidelcartesian__(A,x,y)*2/PR
		ph = 2 * __np__.pi * OPD
		I1 = 1
		I2 = 1
		Ixy = I1 + I2 + 2 * __np__.sqrt(I1*I2) * __np__.cos(ph)
		__tools__.makecircle(Ixy, r, PR)
		fig = __plt__.figure(figsize=(9, 6), dpi=80)
		__plt__.imshow(-Ixy, extent=[-PR,PR,-PR,PR])
		__plt__.set_cmap('Greys')
		__plt__.show()

def __seidelpolar__(coefficient,r,u):
	W = coefficient
	h = 1
	Ap = W[0][0] * h**2
	At = W[1][0] * h*r*__cos__(u-W[1][1]*__np__.pi/180)
	Ad = W[2][0] * r**2
	Aa = W[3][0] * h**2*r**2*__cos__(u-W[3][1]*__np__.pi/180)
	Ac = W[4][0] * h*r*__cos__(u-W[4][1]*__np__.pi/180)
	As = W[5][0] * r**4

	W = Ap+At+Ad+Aa+Ac+As
	return W

def __seidelcartesian__(coefficient,x,y):
	W = coefficient
	h = 1
	u = __arctan2__(y,x)
	r = __sqrt__(x**2+y**2)
	W = __seidelpolar__(coefficient,r,u)
	return W

max_stars_repo_name: Graylien/opticspy

max_issues_repo_name: Graylien/opticspy

max_forks_repo_name: Graylien/opticspy

reason: import numpy,from numpy


text:
"""
Test for all TigerControl environments
"""
import time
import tigercontrol
import jax.numpy as np
import sys
import argparse

def test_environments(stress=True):
    """ Description: test all environments """
    env_ids = tigercontrol.environment_registry.list_ids() # get all controllers

    for env_id in env_ids:
        print("\n------ {} tests -----\n".format(env_id))
        try:
            test_environment(env_id, stress=stress)
        except Exception as e:
            print("ERROR: environment {} raised error {}".format(env_id, e))

def test_environment(env_id, stress=True):
    """ Description: run test for a single environment """
    test_api(env_id)
    if stress:
        test_stress(env_id)

def test_api(env_id):
    """ Description: verify that all default methods work for this environment """
    environment_class = tigercontrol.environment(env_id)
    env = environment_class()
    n, m = env.get_state_dim(), env.get_action_dim() # get dimensions of system
    x_0 = env.reset() # first state
    x_1 = env.step(np.zeros(m)) # try step with 0 action
    assert np.isclose(x_0, env.reset()) # assert reset return back to original state

def test_stress(env_id):
    """ Description: run environment in a few high pressure situations """
    environment_class = tigercontrol.environment(env_id)
    env = environment_class()
    n, m = env.get_state_dim(), env.get_action_dim() # get dimensions of system
    x_0 = env.reset()
    for t in range(10000):
        env.step(t * np.ones(m))

# main file
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--stress", type=bool, help="perform stress test", default=True)
    parser.add_argument("--env", type=str, help="environment id/name", default="all")
    args = parser.parse_args()

    # if no env specified, run tests for all envs
    if args.env == "all":
        test_environments(stress=args.stress)
    # check specified env is in registry and run individual tests
    else:
        if args.env not in tigercontrol.environment_registry.list_ids():
            raise Exception("{} is not a valid env id".format(args.env))
        test_environment(args.env, args.stress)

max_stars_repo_name: MinRegret/TigerControl

max_issues_repo_name: johnhallman/ctsb

max_forks_repo_name: MinRegret/tigercontrol

reason: import jax


text:
import time

from compecon.tools import Options_Container, qzordered
from compecon.nonlinear import MCP
from compecon.lcpstep import lcpstep
from compecon.lqmodel import LQmodel
import numpy as np
import scipy as sp
import pandas as pd
from scipy.sparse import block_diag, kron, issparse, identity
from scipy.sparse.linalg import spsolve
from compecon.tools import jacobian, hessian, gridmake, indices
from inspect import getargspec
#from .lcpstep import lcpstep  # todo: is it worth to add lcpstep?
import warnings

#fixme In docstrings, indicate that discrete model should not include x in definitions

__author__ = 'Randall'

class DPtime(Options_Container):
    def __init__(self, discount=0.0, horizon=np.inf):
        """
        A container class to store time parameters of a DPmodel object

        Parameters
        ----------
        discount : float
            scalar, discount factor in the interval (0, 1)
        horizon: int or np.inf
            agents horizon
        """
        self.discount = discount
        self.horizon = horizon

class DPrandom(Options_Container):
    """ Container for the random components of a DPmodel object

    Attributes:
        e   continuous state transition shocks.
        w   continuous state shock probabilities.
        q   discrete state transition probabilities.
    """
    description = 'Random components of a DPmodel object'

    def __init__(self, ni, nj, e=None, w=None, q=None, h=None):
        if ni == 1:
            qq = np.ones((nj, 1, 1))
        else:
            txt = 'If the model has 2 or more discrete states, a state transition must be provided: ' + \
                  "either deterministic (option 'h') or stochastic ('q'). "
            assert (q is None) ^ (h is None), txt

            qq = np.zeros((nj, ni, ni))
            if q is None:
                h = np.atleast_2d(h)
                for i in range(ni):
                    for j in range(nj):
                        qq[j, i, h[j, i]] = 1
            else:
                q = np.array(q)
                if q.ndim == 2:  # assume it is the same regardless of choice
                    qq[:] = q
                else:
                    qq = q
                ss = qq.shape
                assert ss[1] == ss[2], 'A Markov transition matrix must be square (last two dimensions of q must be equal)'
                # assert np.allclose(qq.sum(-1), 1.0), 'The rows of a Markov transition must add up to 1.0'
                #fixme set tolerance level, seems to be too strict

        self.e = np.zeros((2, 1)) if e is None else np.atleast_2d(e)
        self.w = np.ones((1)) if w is None else np.atleast_1d(w)
        self.q = qq

class DPdims(Options_Container):
    """ Container for the dimensions of a DPmodel object

    Attributes:
        ds  dimension of continuous state s
        dx  dimension of continuous action x
        ni  number of discrete states i
        nj  number of discrete actions j
        ns  number of continuous state nodes output
        nx  number of discretized continuous actions, if applicable
        nc  number of collocation polynomials = of coefficients to be determined
        dr  number of restrictions on x (other than bounds on individual x variables)
    """
    description = "Dimensions of a DPmodel object"

    def __init__(self, ds=1, ns=None,
                 dx=1, nx=None,
                 ni=1, nj=1, ne=1,
                 nc=None):
        self.ds = ds
        self.ns = ns
        self.dx = dx
        self.nx = nx
        self.ni = ni
        self.nj = nj
        self.ne = ne
        self.nc = nc
        self.dr = 0

class DPlabels(Options_Container):
    """ Container for labels of the DPmodel variables

    Attributes:
        s  labels for continuous state variables
        x  labels for continuous action variables
        i  labels for discrete states
        j  labels for discrete choices
    """
    description = "Labels for DPmodel variables"

    def __init__(self, s, x, i, j):
        self.s = s
        self.x = x
        self.i = i
        self.j = j

class DPoptions(Options_Container):
    """ Container for numerical options to solve a DPmodel

    Attributes:
        algorithm             algorithm for solver
        tol                   convergence tolerance parameter
        ncpmethod             method for complementarity problem
        maxit                 maximum number of iterations
        maxitncp              maximunm number of iterations for ncpmethod
        discretized           true if continuous action is discretized or not present.
        X                     dx.nx discretized continuous actions
        D_reward_provided     true if Jacobian and Hessian of reward are provided
        D_transition_provided true if Jacobian and Hessian of transition are provided
        knownFunctions        ni.nj boolean array, true if discrete policy and value functions are known
        print                whether to print output
    """
    description = "Solver options for a DPmodel"

    def __init__(self, algorithm='newton', tol=np.sqrt(np.spacing(1)), ncpmethod='minmax',
                 maxit=80, maxitncp=50, discretized=False, X=None,
                 knownFunctions=None, show=True, print=None):
        self.algorithm = algorithm
        self.tol = tol
        self.ncpmethod = ncpmethod
        self.maxit = maxit
        self.maxitncp = maxitncp
        self.show = show
        self.discretized = discretized
        self.X = X
        self.knownFunctions = knownFunctions
        if print is not None:
            warnings.warn("Keyword 'print=' is deprecated. Use 'show=' instead.")

    def print_header(self, method, horizon):
        horizon = 'infinite' if np.isinf(horizon) else 'finite'
        if self.show:
            print('Solving %s-horizon model collocation equation by %s method' % (horizon, method))
            print('{:4s} {:12s} {:8s}'.format('iter', 'change', 'time'))
            print('-' * 30)

    def print_current_iteration(self, it, change, tic):
        """ Prints summary of current iteration in solve method

        Args:
          it: iteration number (scalar)
          change: distance between two iterations
          tic: time when iterations started

        Returns:
          prints output to screen
        """
        if self.show:
            print('{:4d}  {:12.1e}  {:8.4f}'.format(it, change, time.time() - tic))

    def print_last_iteration(self, tic, change):
        """ Prints summary of last iteration in solve method

        Args:
          tic: time when iterations started
          change: distance between last two iterations

        Returns:
          prints output to screen
        """
        if self.show:
            if change >= self.tol:
                print('Failure to converge in DPmodel.solve()')
            print('Elapsed Time = {:7.2f} Seconds'.format(time.time() - tic))

class DPmodel(object):
    """
        A Dynamic Programming Model class

        A DPmodel object has the following attributes:

        -- Time dimension:
        * horizon:  time horizon (infinite)
        * discount: discount factor (required)

        -- Dimension of state and action spaces:
        * ds: number of continuous state variables (1)
        * dx: number of continuous action variables (1)
        * ni: number of discrete states (1)
        * nj: number of discrete actions (1)

        -- Stochastic components: Markov transition and continuous iid shocks:
        * e:  ne.de discretized continuous state transition shocks (0)
        * w:  ne.1 continuous state shock probabilities (1)
        * q:  ni.ni.nj discrete state transition probabilities (empty)
        * h:  nj.ni deterministic discrete state transitions (empty)
        * X:  nx.dx discretized continuous actions (empty)

        -- Value and policy functions: Interpolator objects:
        * Value:    ni-array for the value function
        * Policy:   ni-array for the policy function
        * Value_j:  nj.ni-array for the value function at each discrete action
        * Policy_j: nj.ni-array for the policy function at each discrete action
        * DiscreteAction: ni.ns-array(integer), discrete actions at each state node

        -- Numerical solution:
        * algorithm:  algorithm for solver
        * tol:        convergence tolerance parameter
        * ncpmethod:  method for complementarity problem
        * maxit:      maximum number of iterations
        * maxitncp:   maximun number of iterations for ncpmethod
        * knownFunctions:  nj.ni-array(boolean), True if policy and value functions are known
        * D_reward_provided: True if derivatives of reward function are provided
        * D_transition_provided: True if derivatives of transition function are provided

        -- Output details and other parameters:
        * nr:      number of refined nodes
        * output:  print output per iterations
        * nc:      number of continuous nodes
        * ns:      number of continuous nodes on output ##TODO
        * xnames:  1.dx cell, cell of names for continuous actions
        * discretized: True if continuous action is discretized or not present.
     """
    # todo: Review the dimensions of attributes in above docstring

    def __init__(self, basis, reward, transition, bounds=None, restrictions=None,
                 i=('State 0',), j=('Choice 0', ), x=(),
                 discount=0.0, horizon=np.inf,
                 e=None, w=None, q=None, h=None, params=None):

        assert callable(reward), 'reward must be a function'
        assert callable(transition), 'transition must be a function'
        if bounds is None:
            assert len(x) == 0, 'Bounds for continuous action are missing!'
        else:
            assert callable(bounds), 'bounds must be a function'

        self.__b = bounds
        self.__f = reward
        self.__g = transition
        self.__h = restrictions

        assert isinstance(i, (list, tuple)), 'i must be a tuple of strings (names of discrete states)'
        assert isinstance(j, (list, tuple)), 'j must be a tuple of strings (names of discrete choices)'
        assert isinstance(x, (list, tuple)), 'x must be a tuple of strings (names of continuous actions)'

        i, j, x = tuple(i), tuple(j), tuple(x)
        ni, nj, dx = len(i), len(j), len(x)

        assert (nj > 1) or (dx > 0), 'Model did not specified any policy variable! Set j or x (or both).'

        #  Value and policy functions
        if np.isinf(horizon):
            self.Value = basis.duplicate(l=[i])
            self.Value_j = basis.duplicate(l=[i, j])
            self.Policy = basis.duplicate(l=[i, x])
            self.Policy_j = basis.duplicate(l=[i, j, x])
            self.DiscreteAction = np.zeros([ni, basis.N], int)
        else:
            t0 = np.arange(horizon)
            t1 = np.arange(horizon + 1)
            self.Value = basis.duplicate(l=[t1, i])
            self.Value_j = basis.duplicate(l=[t1, i, j])
            self.Policy = basis.duplicate(l=[t0, i, x])
            self.Policy_j = basis.duplicate(l=[t0, i, j, x])
            self.DiscreteAction = np.zeros([horizon, ni, basis.N], int)

        # Time parameters
        self.time = DPtime(discount, horizon)

        # Labels for model variables
        self.labels = DPlabels(basis.opts.labels, x, i, j)

        # Stochastic specification
        self.random = DPrandom(ni, nj, e, w, q, h)

        # Default numerical solution parameters and parameters for model functions
        self.options = DPoptions()
        self.params = params

        # Model dimensions
        self.dims = DPdims(basis.d,  # number of continuous state variables
                           basis.N,  # number of continuous state nodes
                           dx,  # number of continuous policy variables
                           0,  # number of discretized policy values
                           ni,  # number of discrete states
                           nj,  # number of discrete choices
                           self.random.w.size,  # number of discretized continuous shocks
                           basis.M)  # number of collocation coefficients

        ''' <<<<<<<<<<<<<<<<<<<             END OF CONSTRUCTOR        >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>'''

    def __repr__(self):
        txt = 'A continuous state, ' + ('continuous' if self.dims.dx > 0 else 'discrete') + ' action dynamic model.\n'
        txt = txt.upper()
        txt += '\n\t* Continuous states:\n'
        n, a, b = self.Value.n, self.Value.a, self.Value.b
        for k in range(self.Value.d):
            txt += "\t\t{:<2d}:  {:<s} --> {:d} nodes in [{:.2f}, {:.2f}]\n".format(k, self.labels.s[k], n[k], a[k], b[k])
        if self.dims.dx > 0:
            txt += '\n\t* Continuous actions\n'
            for v, vlab in enumerate(self.labels.x):
                txt += '\t\t{:<2d}:  {:s}\n'.format(v, vlab)
        if self.dims.ni > 1:
            txt += '\n\t* Discrete states\n'
            for v, vlab in enumerate(self.labels.i):
                txt += '\t\t{:<2d}:  {:s}\n'.format(v, vlab)
        if self.dims.nj > 1:
            txt += '\n\t* Discrete choices:\n'
            for v, vlab in enumerate(self.labels.j):
                txt += '\t\t{:<2d}:  {:s}\n'.format(v, vlab)

        return txt

    def bounds(self, s, i, j):  # --> (lowerBound, UpperBound)
        """ Returns upper-  and lower-bounds for the continuous action variable.

        Depends only on state variables and discrete action
        """
        ns = s.shape[-1]
        dx, ds = self.dims['dx', 'ds']
        lb, ub = self.__b(s, i, j) #, *self.__par['bounds']) only if params defined
        lb.shape = dx, ns
        ub.shape = dx, ns
        return lb, ub

    def restrictions(self, s, x, i, j, derivative=False):  # --> (h, hx, hxx)
        """ Returns the restriction functions and their first- and second-derivatives.

        Depends only on current variables

        """
        ns = s.shape[-1]
        dx = self.dims.dx

        hh = self.__h(s, x, i, j)

        if isinstance(hh, tuple):
            # assert len(ff) == 3, 'reward must return 1 or 3 arrays'  # commented-out for speed
            h, hx, hxx = hh[0].reshape(-1, ns), hh[1].reshape(-1, dx, ns), hh[2].reshape(-1, dx, dx, ns)
        else:
            h, hx, hxx = hh.reshape(-1, ns), None, None

        return (h, hx, hxx) if derivative else h

    def reward(self, s, x, i, j, derivative=False):  # --> (f, fx, fxx)
        """ Returns the reward function (e.g. utility, profits) and its first- and second-derivatives.

        Depends only on current variables

        """
        ns = s.shape[-1]
        dx = self.dims.dx

        ff = self.__f(s, x, i, j)
        if isinstance(ff, tuple):
            # assert len(ff) == 3, 'reward must return 1 or 3 arrays'  # commented-out for speed
            f, fx, fxx = ff[0].reshape(1, ns), ff[1].reshape(dx, ns), ff[2].reshape(dx, dx, ns)
        else:
            f, fx, fxx = ff.reshape(1, ns), None, None

        return (f, fx, fxx) if derivative else f

    def transition(self, s, x, i, j, in_, e, derivative=False):  # --> (g, gx, gxx)
        """ Returns the next-period continuous state and its first- and second-derivatives.

         Depends on current (s,x,i,j) and future (in,e) variables
        """
        ns = s.shape[-1]
        dx, ds = self.dims['dx', 'ds']

        gg = self.__g(s, x, i, j, in_, e)
        if isinstance(gg, tuple):
            # assert len(ff) == 3, 'reward must return 1 or 3 arrays'  # commented-out for speed
            g, gx, gxx = gg[0].reshape(ds, ns), gg[1].reshape(dx, ds, ns), gg[2].reshape(dx, dx, ds, ns)
        else:
            g, gx, gxx = gg.reshape(ds, ns), None, None

        return (g, gx, gxx) if derivative else g

    def solve(self, v=None, x=None, nr=10, **kwargs):
        """ Solves the model

        Args:
          solver: a DPsolver object (optional)

        Returns:

        """

        # Update solution options using kwargs
        self.options[kwargs.keys()] = kwargs.values()

        t = slice(None) if np.isinf(self.time.horizon) else -1  # if finite horizon, v is taken as last period

        if v is not None:
            self.Value[t] = v[t]

        if x is not None:
            self.Policy_j[t] = x[t]

        ''' 1: PREPARATIONS*********************** '''
        ni, nj, dx = self.dims['ni', 'nj', 'dx']

        if self.options.knownFunctions is None:
            self.options.knownFunctions = np.zeros([ni, nj], bool)

        # Step necessary for discretized models
        if dx == 0:
            self.options.discretized = True
        elif self.options.X is not None:
            self.options.X = np.atleast_2d(self.options.X)
            self.dims.nx = self.options.X.shape[-1]  # number of discretized policy values
            self.options.discretized = True
            if self.options.X.shape[0] != dx:
                raise ValueError('If model is discretized, field "X" must have {} rows'.format(dx))

        # s0 = self.Value.nodes
        # x0 = self.Policy.y[..., 0, :, :]
        # self.options.D_reward_provided = isinstance(self.__f(s0, x0, 0, 0), tuple)
        # self.options.D_transition_provided = isinstance(self.__g(s0, x0, 0, 0, 0, 0), tuple)

        ''' 2: SOLVE THE MODEL******************** '''
        if np.isfinite(self.time.horizon):
            self.__solve_backwards()
        elif self.options.algorithm == 'funcit':
            self.__solve_by_function_iteration()
        elif self.options.algorithm == 'newton':
            self.__solve_by_Newton_method()
        else:
            raise ValueError('Unknown solution algorithm')

        self.update_policy()

        if nr is not None:
            return self.solution(nr)

    def solution(self, nr=10, resid=True):
        """
        Computes solution over a refined grid

        nr:  scalar or np.array
            -- scalar >> compute a grid from basis
            -- array  >> compute solution over provided values
        resid: compute residuals if True

        """
        # TODO:  Make finite horizon case

        ni, nj, ds, dx = self.dims['ni', 'nj','ds', 'dx']

        scalar_input = np.isscalar(nr) and isinstance(nr, int)

        if scalar_input:
            a = self.Value.a
            b = self.Value.b
            n = self.Value.n
            sr = np.atleast_2d(gridmake(*[np.linspace(a[i], b[i], nr * n[i]) for i in range(self.Value.d)]))
        else:
            sr = np.atleast_2d(nr)
            assert sr.shape[0] == self.dims.ds, 'provided s grid must have {} rows'.format(self.dims.ds)

        ''' MAKE DATABASE'''
        # ADD CONTINUOUS STATE VARIABLE
        DATA = pd.DataFrame(np.tile(sr, ni).T, columns=self.labels.s)

        # ADD DISCRETE STATE VARIABLE
        if ni > 1:
            ivals= np.repeat(np.arange(ni), sr.shape[1])
            icat = pd.Categorical.from_codes(ivals, self.labels.i)
            DATA['i'] = ivals

        # SET INDEX FOR DATA
        if ds == 1:
            slab = DATA[self.labels.s[0]]
            if ni > 1:
                DATA.index = pd.MultiIndex.from_arrays([icat, slab])
            else:
                DATA.index = slab
        elif ni > 1:
            DATA.index = icat

        # COMPUTE OPTIMAL POLICY AND VALUE
        xr = self.Policy_j(sr, dropdim=False)  # [0] because there is only 1 order
        vr = self.vmax(sr, xr, self.Value)
        v_LHS = self.Value(sr, dropdim=False) # LHS of Bellman equation: V(s)

        # ADD VALUE FUNCTION
        DATA['value'] = v_LHS.flatten()

        # ADD RESIDUAL IF REQUESTED
        if resid:
            v_RHS = np.max(vr, -2)  # RHS of Bellman equation: max[f(s,x) + dV(s')]
            DATA['resid'] = (v_LHS - v_RHS).flatten()

        # ADD VALUE FUNCTION PER DISCRETE ACTION
        if nj > 1:
            tempj = np.argmax(vr,-2).flatten()
            DATA['j*'] = pd.Categorical.from_codes(tempj,self.labels.j)  # optimal discrete choice

            for j, jlabel in enumerate(self.labels.j):
                DATA['value[' + jlabel + ']'] = vr[:, j].flatten()

        # ADD CONTINUOUS ACTION
        if dx:
            for ix, xlabel in enumerate(self.labels.x):
                DATA[xlabel] = self.Policy(sr, dropdim=False)[:,ix].flatten()

                # ADD CONTINUOUS ACTION PER DISCRETE ACTION
                if nj > 1:
                    for j, jlabel in enumerate(self.labels.j):
                        DATA[xlabel + '[' + jlabel + ']'] = xr[:, j, ix].flatten()

        return DATA

    def simulate(self, nper, sinit, iinit=0, seed=None):

        # Simulate the model
        #
        #   S = self.simulate(nper, sinit, iinit)
        #
        # nper = number of periods to simulate (scalar)
        # sinit = initial continuos state (nrep x ds), where nrep is number of repetitions
        # iinit = initial discrete state (scalar)
        #
        # S = simulation results (table), with variables:
        #    r, repetion number
        #    t, time period
        #    i, discrete state
        #    j, optimal discrete action
        #    s, continuous state
        #    x, optimal continuous action

        # ****** 1: Preparation***********************************************************
        ds, dx, ni, nj, ne = self.dims['ds', 'dx', 'ni', 'nj', 'ne']

        assert iinit < ni, 'Invalid initial discrete state: must be an integer between 0 and {}'.format(ni - 1)
        if seed:
            np.random.seed(seed)

        # Determine number of replications nrep and periods nper to be simulated.
        # nper cannot exceed time.horizon.
        sinit = np.atleast_2d(sinit).astype(float)
        ds2, nrep = sinit.shape
        assert ds==ds2, 'initial continous state must have {} rows'.format(ds)
        nper = min(nper, self.time.horizon) + 1

        ### Allocate memory to output arrays
        ssim = np.empty((nper + 1, ds, nrep))
        xsim = np.empty((nper, dx, nrep))
        isim = np.empty((nper + 1, nrep), dtype=int)
        jsim = np.empty((nper, nrep), dtype=int)

        ### Set initial states
        ss = sinit
        ii = np.full(nrep, iinit)
        ssim[0] = ss
        isim[0] = ii

        # ***** *2: Simulate the model ***************************************************
        #
        # For each period
        for ip in range(nper):
            # # *TODO* Write the finite horizon case
            # #{
            # if T<inf
            #     for i=1:ni
            #         for j=1:nj
            #             cx[:,:,i,j] = Phi\x[:,:,i,j,ip]
            #             cv[:,i,j]   = Phi\v[:,i,j,ip]
            #         end
            #     end
            # end
            #}

            ### Allocate memory for current-period policy/value functions, all repetitions
            xx = np.empty_like(xsim[0])
            jj = np.empty_like(jsim[0])
            vv = np.empty((nj, nrep))

            ### For the current discrete state ii, compute the conditional value functions use this to determine the
            # optimal discrete policy jmax
            for i in range(ni):
                ir = ii == i
                if not np.any(ir):
                    continue
                vv[:, ir] = self.Value_j[i](ss[:, ir], dropdim=False)
            jmax = np.argmax(vv, 0)

            ### For current discrete state/policy,  interpolate the optimal continuous policy, keeping it within bounds
            #  record discrete policy
            for i in range(ni):
                for j in range(nj):
                    ir = (ii == i) & (jmax == j)
                    if not np.any(ir):
                        continue
                    if dx > 0:
                        xx[:, ir] = self.Policy_j[i, j, :](ss[:, ir])
                        xl, xu = self.bounds(ss[:, ir], i, j)
                        xx[:, ir] = np.minimum(np.maximum(xx[:, ir], xl), xu)

                    jj[ir] = j

            ### Compute the new continuous shock (necessary for the transition function)
            if self.random.w.size == 1:
                ee = np.tile(self.random.e, nrep)     #self.random.e[ones(nrep,1),:]  #fixme make sure dimensions are ok
            else:
                ee = self.random.e[:, np.random.choice(ne, nrep, p=self.random.w)] #fixme make sure dimensions are ok

            ### Compute the new state: use the Markov transition matrix self.random.q to update
            # the discrete state then update the continuous state
            iiold = ii[:]
            ii = np.zeros(nrep)
            for i in range(ni):
                for j in range(nj):
                    ir = (iiold == i) & (jj == j)
                    if not np.any(ir):
                        continue
                    if ni > 1:
                        ii[ir] =  np.random.choice(ni, ir.sum(), p=self.random.q[j, i, :])
                        # ii[ir] =  discrand(ir.sum(), self.random.q[j, i, :])
                    ss[:, ir] = self.transition(ss[:, ir], xx[:, ir], i, j, ii[ir], ee[:, ir])

            ### Save the current-period simulation
            ssim[ip + 1] = ss
            isim[ip + 1] = ii
            xsim[ip] = xx
            jsim[ip] = jj

        ### Trim the last observation if model has infinite horizon
        if np.isinf(self.time.horizon):
            ssim = ssim[:nper]
            isim = isim[:nper]

        # ****** 3: Make a table with the simulated data *********************************
        #
        # # Permute (transpose) the dimensions of the data, so first index refer to time.
        # ssim = permute(ssim,[2,1,3])
        # xsim = permute(xsim,[2,1,3])
        # isim = isim'
        # jsim = jsim'

        ### Add variables rsim and tsim to identify the repetition number and the time
        # period of each observation
        tsim, rsim = gridmake(np.arange(nper), np.arange(nrep))

        ### Make the table.
        DATA = pd.DataFrame()
        DATA['time'] = tsim
        if nrep > 1:
            DATA['_rep'] = rsim

        if ni > 1:
            DATA['i'] = pd.Categorical.from_codes(isim.flatten(),self.labels.i)

        sdata = ssim.swapaxes(0, 1).reshape((ds, -1))
        for k, slab in enumerate(self.labels.s):
            DATA[slab] = sdata[k]

        if nj > 1:
            DATA['j*'] = pd.Categorical.from_codes(jsim.flatten(), self.labels.j)

        if dx > 0:
            xdata = xsim.swapaxes(0, 1).reshape((dx, -1))
            for k, xlab in enumerate(self.labels.x):
                DATA[xlab] = xdata[k]

        return DATA

    def lqapprox(self, s0, x0):
        """
        Solves discrete time continuous state/action dynamic programming model using a linear quadratic approximation
        Args:
            s0: steady-state state
            x0:  steady-state action

        Returns:
            A LQmodel object
        """

        assert (self.dims.ni * self.dims.nj < 2), 'Linear-Quadratic not implemented for models with discrete state or choice'
        s0, x0 = np.atleast_1d(s0, x0)

        assert s0.size == self.dims.ds, 's0 must have %d values' % self.dims.ds
        assert x0.size == self.dims.dx, 'x0 must have %d values' % self.dims.dx

        s0, x0 = s0.astype(float), x0.astype(float)
        s0.shape = -1, 1
        x0.shape = -1, 1

        delta = self.time.discount

        # Fix shock at mean
        estar = self.random.w @ self.random.e.T
        estar.shape = -1, 1

        # Get derivatives
        f0, fx, fxx = self.reward(s0, x0, None, None, True)
        g0, gx, gxx = self.transition(s0, x0, None, None, None, estar, derivative=True)

        fs = jacobian(lambda y: self.reward(y.reshape((-1, 1)), x0, None, None), s0)
        fxs = jacobian(lambda y: self.reward(y.reshape((-1, 1)), x0, None, None, True)[1], s0)
        fss = hessian(lambda y: self.reward(y.reshape((-1, 1)), x0, None, None), s0)
        gs = jacobian(lambda y: self.transition(y.reshape((-1, 1)), x0, None, None, None, estar), s0)

        # Reshape to ensure conformability
        ds, dx = self.dims['ds', 'dx']

        f0.shape = 1, 1
        s0.shape = ds, 1
        x0.shape = dx, 1
        fs.shape = 1, ds
        fx.shape = 1, dx
        fss.shape = ds, ds
        fxs.shape = dx, ds
        fxx.shape = dx, dx
        g0.shape = ds, 1
        gx.shape = ds, dx
        gs.shape = ds, ds
        fsx = fxs.T

        f0 += - fs @ s0 - fx @ x0 + 0.5 * s0.T @ fss @ s0 + s0.T @ fsx @ x0 + 0.5 * x0.T @ fxx @ x0
        fs += - s0.T @ fss - x0.T @ fxs
        fx += - s0.T @ fsx - x0.T @ fxx
        g0 += - gs @ s0 - gx @ x0

        return LQmodel(f0, fs, fx, fss, fsx, fxx, g0, gs, gx, delta,self.labels.s, self.labels.x)

    def __solve_backwards(self):
        """
        Solve collocation equations for finite horizon model by backward recursion
        """
        T = self.time.horizon
        s = self.Value.nodes

        tic = time.time()
        self.options.print_header('backward recursion', T)
        for t in reversed(range(T)):
            self.options.print_current_iteration(t, 0, tic)
            self.Value_j[t] = self.vmax(s,
                                        self.Policy_j.y[t],
                                        self.Value[t + 1])
            self.make_discrete_choice(t)

        self.options.print_last_iteration(tic, 0)
        return None

    def __solve_by_function_iteration(self):
        """
            Solves infinite-horizon model collocation equation by function iteration. Solution is found when the
            collocation coefficients of the value function converge to a fixed point (within |self.tol| tolerance).
         """
        tic = time.time()
        s = self.Value.nodes
        self.options.print_header('function iteration', self.time.horizon)
        for it in range(self.options.maxit):
            cold = self.Value.c.copy()
            self.Value_j[:] = self.vmax(s, self.Policy_j.y, self.Value)
            self.make_discrete_choice()
            change = np.linalg.norm((self.Value.c - cold).flatten(), np.Inf)
            self.options.print_current_iteration(it, change, tic)
            if change < self.options.tol:
                break
            if np.isnan(change):
                raise ValueError('nan found on function iteration')
        self.options.print_last_iteration(tic, change)

    def __solve_by_Newton_method(self):
        tic = time.time()
        s = self.Value_j.nodes
        x = self.Policy_j.y
        ni = self.dims.ni

        if issparse(self.Value_j._Phi):
            Phik = kron(identity(ni), self.Value_j._Phi, format='csr').toarray()
            # Solve = spsolve
        else:
            Phik = np.kron(np.eye(self.dims.ni), self.Value._Phi)
            # Solve = np.linalg.solve

        # todo: fix the dimensions and check that Phik is transposed?

        def SOLVE(A,b):
            return np.linalg.solve(A, b) if (self.dims.ns == self.dims.nc) else np.linalg.lstsq(A, b)[0]

        self.options.print_header("Newton's", self.time.horizon)
        for it in range(self.options.maxit):
            cold = self.Value.c.copy().flatten()
            # print('\ncold', cold)
            self.Value_j[:], vc = self.vmax(s, x, self.Value, True)
            self.make_discrete_choice()
            step = - SOLVE(Phik - vc, Phik @ cold - self.Value.y.flatten())
            c = cold + step
            change = np.linalg.norm(step, np.Inf)
            self.Value.c = c.reshape(self.Value.c.shape)
            self.options.print_current_iteration(it, change, tic)
            if np.isnan(change):
                raise ValueError('nan found on Newton iteration')
            if change < self.options.tol:
                break
        self.options.print_last_iteration(tic, change)

    def vmax(self, s, x, Value, dVc=False):  # [v,x,vc]
        # Unpack model structure
        ni, nj = self.dims['ni', 'nj']
        ns = s.shape[-1]
        ms = Value.M  # number of polynomials
        v = np.empty([ni, nj, ns])

        if self.dims.dx == 0:  # Discrete model
            # hh = slice(None)
            for i in range(ni):
                for j in range(nj):
                    v[i, j] = self.__Bellman_rhs_discrete(Value, None, s, i, j)
        elif self.options.discretized:
            # hh = slice(None)
            for i in range(ni):
                for j in range(nj):
                    v[i, j] = self.vmax_discretized(Value, s, x[i, j], i, j)
        elif self.__h is None:
            # hh = 0
            for i in range(ni):
                for j in range(nj):
                    v[i, j] = self.vmax_continuous(Value, s, x[i, j], i, j)
        else:
            for i in range(ni):
                for j in range(nj):
                    v[i, j] = self.vmax_continuous_restricted(Value, s, x[i, j], i, j)

        if not dVc:
            return v

        # Computes derivative with respect to Value function interpolation coefficients
        e, w, q = self.random['e', 'w', 'q']

        if ni * nj > 1:
            vc = np.zeros((ns, ni, ms, ni))
            jmax = np.argmax(v, 1)

            for i in range(ni):
                for j in range(nj):
                    is_ = jmax[i] == j
                    if not np.any(is_):
                        continue

                    for k in range(w.size):
                        ee = np.tile(e[:, [k]], ns)  # indexing with [k] instead of k retains shape of vector!
                        for in_ in range(ni):
                            if q[j, i, in_] > 0:
                                snext = self.transition(s[:, is_], x[i, j, :, is_], i , j, in_, ee[:, is_])  #fixme need to know number of output arguments!!!
                                prob = w[k] * q[j, i, in_,]
                                vc[is_, i, :, in_] += prob * Value.Phi(snext).toarray().reshape((is_.sum(), ms), order='F')   #fixme I can't find the proper way to index this

            vc = vc.reshape((ns*ni,ms*ni),order='F')
        else:
            vc = np.zeros((ns, ms))
            for k in range(w.size):
                ee = np.tile(e[:, [k]], ns)
                snext = self.transition(s, x[0, 0], 0, 0, 0, ee) #fixme need to know number of output arguments!!!
                vc += w[k] * Value.Phi(snext)

        vc *= self.time.discount
        return v, vc

    # vmax_discretized
    # Nested function in vmax: Finds the optimal policy and value function for a given pair of discrete state
    # and discrete action, when the continuous policy has been discretized.

    def vmax_discretized(self, Value, s, xij, i, j):

        nx = self.dims.nx
        ns = s.shape[-1]
        dx = self.dims.dx
        X = self.options.X

        vv = np.full((nx, ns), -np.inf)

        xl, xu = self.bounds(s, i, j)
        xl = xl.T
        xu = xu.T

        for h, x0 in enumerate(X.T):
            is_= np.all((xl <= x0) & (x0 <= xu), 1)
            if np.any(is_):
                xx = np.repeat(x0, ns, 0)
                vv[h, is_] = self.__Bellman_rhs_discrete(Value, xx, s, i, j)

        xmax = np.argmax(vv, 0)

        vxs = [a[0] for a in np.indices(vv.shape)]  # the [0] reduces one dimension
        vxs[0] = xmax
        vij = vv[vxs]

        xxs = [a[0] for a in np.indices(X.T.shape)]
        xxs[0] = xmax
        xij[:] = X.T[xxs]

        return vij

    # vmax_continuous
    # Nested function in vmax: Finds the optimal policy and value function for a given pair of discrete state
    # and discrete action, by solving the linear complementarity problem.
    def vmax_continuous_MCP(self, Value, s, xij, i, j):

        ns = s.shape[-1]
        dx = self.dims.dx

        def kkt(xvec):
            """ Karush-Kuhn Tucker conditions

            The first-order conditions are given by the derivative of Bellman equation. The problem needs to be
            solved ns times (number of nodes), each time with dx unknowns (policy variables). This routine
            expresses all this FOCs as a single vector, making a block-diagonal matrix with the respective Hessian matrices
            (= jacobian of FOCs). The resulting output is suitable to be solved by the MCP class.

            """
            xij = xvec.reshape((dx, ns))
            EV, EVx, EVxx = self.__Bellman_rhs(Value, s, xij, i, j)

            # and let the first index indicate node
            Vx = np.swapaxes(EVx, 0, -1)
            Vxx = np.swapaxes(EVxx, 0, -1)
            return Vx.flatten(), block_diag(Vxx, 'csc')#.toarray()  #todo not so sure I want a full array, but a lot of trouble with sparse

        xl, xu = self.bounds(s, i, j)

        xlv, xuv, xijv = map(lambda z: z.flatten('F'), (xl, xu, xij))  # vectorize

        FOCs = MCP(kkt, xlv, xuv, xijv)
        xij[:] = FOCs.zero(show=False, transform=self.options.ncpmethod).reshape((ns, dx)).T

        return self.__Bellman_rhs(Value, s, xij, i, j)[0][0]

    # Nested function in vmax: Finds the optimal policy and value function for a given pair of discrete state
    # and discrete action, by solving the linear complementarity problem.
    def vmax_continuous(self, Value, s, xij, i, j):
        xl, xu = self.bounds(s, i, j)

        for it in range(self.options.maxitncp):
            vv, vx, vxx = self.__Bellman_rhs(Value, s, xij, i, j)

            # Compute Newton step, update continuous action, check convergence
            vx, delx = lcpstep(self.options.ncpmethod, xij, xl, xu, vx, vxx)
            xij[:] += delx
            if np.linalg.norm(vx.flatten(), np.Inf) < self.options.tol:
                break

        return self.__Bellman_rhs(Value, s, xij, i, j)[0][0]

    def vmax_continuous_restricted(self, Value, s, xij, i, j):

        ns = s.shape[-1]
        dx = self.dims.dx
        xl, xu = self.bounds(s, i, j)

        dr = self.restrictions(s, xij, i, j).shape[0]
        lij = np.ones((dr, ns))
        xl = np.vstack((xl, np.zeros((dr, ns))))
        xu = np.vstack((xu, np.tile(np.inf, (dr, ns))))
        ZEROS = np.zeros((dr, dr, ns))

        for it in range(self.options.maxitncp):
            vv, vx, vxx = self.__Bellman_rhs(Value, s, xij, i, j)

            # Adjust derivative to add the Lagrange multiplier term
            h, hx, hxx = self.restrictions(s, xij, i, j, True)

            for ir in range(dr):
                vx += lij[ir] * hx[ir]
                vxx += lij[ir] * hxx[ir]

            vx = np.vstack((vx, -h))

            vxx = np.vstack((np.hstack((vxx, np.swapaxes(hx, 0,1))),
                             np.hstack((-hx, ZEROS))))

            #  Compute Newton step, update continuous action, check convergence
            vx, delxl = lcpstep(self.options.ncpmethod,
                               np.vstack((xij, lij)),
                               xl, xu, vx, vxx)
            delx, dell = np.split(delxl, [dx])
            xij[:] += delx
            lij[:] += dell

            print('it = ', it, '\tchange = ', np.linalg.norm(vx.flatten(), np.Inf))
            if np.linalg.norm(vx.flatten(), np.Inf) < self.options.tol:
                break

        return self.__Bellman_rhs(Value, s, xij, i, j)[0][0]

    def __Bellman_rhs_discrete(self, Value, xij, s, i, j):
        ni, nj = self.dims['ni', 'nj']
        ns = s.shape[-1]
        e, w, q = self.random['e', 'w', 'q']
        vv = self.reward(s, xij, i, j)

        for k in range(w.size):
            # Compute states next period
            ee = np.tile(e[:, k: k + 1], ns)
            for in_ in range(ni):
                if q[j, i, in_] == 0:
                    continue
                snext = np.real(self.transition(s, xij, i, j, in_, ee))
                prob_delta = self.time.discount * w[k] * q[j, i, in_]
                vv += prob_delta * Value[in_](snext)
        return vv

    def __Bellman_rhs(self, Value, s, xij, i, j):
        ni, nj = self.dims['ni', 'nj']  # dimensions
        ns = s.shape[-1]
        e, w, q = self.random['e', 'w', 'q']                # randomness

        vv, vx, vxx = self.reward(s, xij, i, j, True)

        for k in range(w.size):
            # Compute states next period and derivatives
            ee = np.tile(e[:, [k]], ns)
            for in_ in range(ni):
                if q[j, i, in_] == 0:
                    continue

                snext, snx, snxx = self.transition(s, xij, i, j, in_, ee, derivative=True)
                snext = np.real(snext)
                prob_delta = self.time.discount * w[k] * q[j, i, in_]

                vn, vns, vnss = Value[in_](snext, order='all', dropdim=False)  # evaluates function, jacobian, and hessian if order='all'

                # Drop the unnecessary dimension
                vns = vns[:, 0]
                vnss = vnss[:, :, 0]

                vv += prob_delta * vn
                vx += prob_delta * np.einsum('k...,jk...->j...', vns, snx)
                vxx += prob_delta * np.einsum('hi...,ij...,kj...->hk...', snx, vnss, snx) +  \
                      np.einsum('k...,ijk...->ij...', vns, snxx)

        return vv, vx, vxx

    def make_discrete_choice(self, t=None):
        # notice : Value_j.y  dims are: 0=state, 1=action, 2=node

        if self.dims.nj == 1:
            if t is None:
                self.Value[:] = self.Value_j.y[:, 0]
            else:
                self.Value[t] = self.Value_j.y[t]
            return

        if t is None:
            self.DiscreteAction = np.argmax(self.Value_j.y, 1)
            ijs = [a[:, 0] for a in np.indices(self.Value_j.y.shape)]
            ijs[1] = self.DiscreteAction
            self.Value[:] = self.Value_j.y[ijs]

        else:
            self.DiscreteAction[t] = np.argmax(self.Value_j.y[t], 1)
            ijs = [a[:, 0] for a in np.indices(self.Value_j.y[t].shape)]
            ijs[1] = self.DiscreteAction[t]
            self.Value[t] = self.Value_j.y[t][ijs]

    def update_policy(self):
        if self.dims.nj == 1:
            self.Policy[:] = self.Policy_j.y[:, 0]
        else:
            ijxs = [a[:, 0] for a in np.indices(self.Policy_j.y.shape)]
            ijxs[-3] = self.DiscreteAction[:, np.newaxis, :]
            self.Policy[:] = self.Policy_j.y[ijxs]

    def check_derivatives(self):
        ni, nj, ds, ns = self.dims['ni', 'nj', 'ds', 'ns']
        nij = ni * nj
        idx = pd.MultiIndex.from_product([np.arange(ni), np.arange(nj)], names=['i', 'j'])
        F = pd.DataFrame({
            'fx': np.zeros(nij) + np.nan,
            'fxx': np.zeros(nij) + np.nan},
            index=idx)

        idy = pd.MultiIndex.from_product([['gx', 'gxx'], np.arange(ds)], names=['func', 's'])
        G = pd.DataFrame(np.zeros((nij,2 * ds)), index=idx, columns=idy)

        s = self.Value.nodes
        e = np.tile(self.random.e[:, 0], ns)

        for k, (i, j) in enumerate(indices(ni,nj).T):
            xl, xu = self.bounds(s, i , j)
            xlinf = np.isinf(xl)
            xuinf = np.isinf(xu)
            xm = (xl + xu) / 2
            xm[xlinf] = xu[xlinf]
            xm[xuinf] = xl[xuinf]
            xm[xuinf & xlinf] = 0.0

            f, fx, fxx = self.reward(s,xm,i,j, derivative=True)
            fxa = jacobian(lambda z: self.reward(s, z, i, j), xm)
            fxxa = hessian(lambda z: self.reward(s, z, i, j), xm)

            #FIXME: HESSIAN SEEMS TO BE THROWING WRONG DIMENSIONS!!

            F.values[k, 0] = np.linalg.norm(fx - fxa, np.inf)
            F.values[k, 1] = np.linalg.norm(fxx - fxxa, np.inf)

    def __get_derivatives(self):
        pass

# TODO: design this class:
"""
    * Should finite-infinite horizon models be split into two subclasses?
    * Should discretized models be handled by a subclass?
    * Should vmax operate directly on Value_j and Policy_j? how to deal with residuals?
"""

max_stars_repo_name: randall-romero/CompEcon-python

max_issues_repo_name: randall-romero/CompEcon

max_forks_repo_name: randall-romero/CompEcon-python

reason: import numpy,import scipy,from scipy