text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/1/10 下午2:19
# @Author  : FengDa
# @File    : two_stage_flow_test.py
# @Software: PyCharm
import torch
import torch.nn as nn
import numpy as np
import mmcv
import torchvision.models as models
import torch.nn.functional as F

from .base import BaseDetector
from .test_mixins import RPNTestMixin, BBoxTestMixin, MaskTestMixin
from .. import builder
from ..registry import DETECTORS
from mmdet.core import bbox2roi, bbox2result, build_assigner, build_sampler
from mmdet.core import bbox_overlaps, bbox2result_with_id
# from ..flow_heads.propagation_utils import warp
from ..flow_heads.propagation_utils_deprecated import warp
# from ..decision_net.utils import resize, change_img_meta, trans_action
from mmdet.models.flow_heads import FlowNetC, flownetc

@DETECTORS.register_module
class TwoStageDetectorFlowTest(BaseDetector, RPNTestMixin,
                           MaskTestMixin):
    def __init__(self,
                 backbone,
                 neck=None,
                 rpn_head=None,
                 bbox_roi_extractor=None,
                 bbox_head=None,
                 track_head=None,
                 mask_roi_extractor=None,
                 mask_head=None,
                 flow_head=None,
                 train_cfg=None,
                 test_cfg=None,
                 pretrained=None):
        """
        :param backbone: ConfigDict,采用的backbone网络结构
        :param neck: ConfigDict,采用的neck网络结构
        :param rpn_head: ConfigDict,采用的rep_head网络结构
        :param bbox_roi_extractor: ConfigDict,采用的bbox_roi_extractor网络结构
        :param bbox_head: ConfigDict,采用的bbox_head网络结构
        :param track_head: ConfigDict,采用的track_head网络结构
        :param mask_roi_extractor: ConfigDict,采用的mask_roi_extractor网络结构
        :param mask_head: ConfigDict,采用的mask_head网络结构
        :param train_cfg: ConfigDict, 训练配置参数
        :param test_cfg: ConfigDict, 测试配置参数
        :param pretrained: str, e.g.: 'modelzoo://resnet50'
        """

        super(TwoStageDetectorFlowTest, self).__init__()
        self.backbone = builder.build_backbone(backbone)

        if neck is not None:
            self.neck = builder.build_neck(neck)
        else:
            raise NotImplementedError

        if rpn_head is not None:
            self.rpn_head = builder.build_head(rpn_head)

        if bbox_head is not None:
            self.bbox_roi_extractor = builder.build_roi_extractor(
                bbox_roi_extractor)
            self.bbox_head = builder.build_head(bbox_head)
        if track_head is not None:
            self.track_head = builder.build_head(track_head)
        if mask_head is not None:
            self.mask_roi_extractor = builder.build_roi_extractor(
                mask_roi_extractor)
            self.mask_head = builder.build_head(mask_head)
        if flow_head is not None:
            self.flow_head = builder.build_head(flow_head)

        self.flow_head1 = None

        self.train_cfg = train_cfg
        self.test_cfg = test_cfg

        self.prev_bboxes = None
        self.prev_roi_feats = None
        self.init_weights(pretrained=pretrained)

    @property
    def with_rpn(self):
        return hasattr(self, 'rpn_head') and self.rpn_head is not None

    @property
    def with_flow(self):
        return self.flow_head is not None

    def load_flow(self):
        # flow_network_data = torch.load("/home/ubuntu/code/fengda/MaskTrackRCNN/pretrained_models/flownetc_EPE1.766.tar")
        # print("=> using pre-trained model '{}'".format(flow_network_data['arch']))
        # self.flow_model = FlowNetC(batchNorm=False)
        # self.flow_model.load_state_dict(flow_network_data['state_dict'])
        # # flow_model = models.__dict__[flow_network_data['arch']](flow_network_data).cuda()
        # self.flow_model.cuda()
        # self.flow_model.eval()
        # from mmdet.models.flow_heads import FlowNetC
        self.flow_head = FlowNetC(batchNorm=False,
            checkpoint="/home/ubuntu/code/fengda/MaskTrackRCNN/pretrained_models/flownetc_EPE1.766.tar")
        # self.flow_head1.requires_grad = False

    def init_weights(self, pretrained=None):
        """initialize weights."""
        super(TwoStageDetectorFlowTest, self).init_weights(pretrained)
        self.backbone.init_weights(pretrained=pretrained)
        if self.with_neck:
            if isinstance(self.neck, nn.Sequential):
                for m in self.neck:
                    m.init_weights()
            else:
                self.neck.init_weights()
        if self.with_rpn:
            self.rpn_head.init_weights()
        if self.with_bbox:
            self.bbox_roi_extractor.init_weights()
            self.bbox_head.init_weights()
        if self.with_mask:
            self.mask_roi_extractor.init_weights()
            self.mask_head.init_weights()
        if self.with_track:
            self.track_head.init_weights()
        # if self.with_flow:
        #     self.flow_head.init_weights()
        #     self.flow_head.cuda()
        #     self.flow_head.requires_grad = False

    def extract_feat(self, img):
        """extract feature map with backbone and neck."""
        x, feat_res0 = self.backbone(img)
        # feat_resnet = x
        if self.with_neck:
            x = self.neck(x)
        return x, feat_res0

    def forward_train(self,
                      img,
                      img_meta,
                      gt_bboxes,
                      gt_bboxes_ignore,
                      gt_labels,
                      ref_img,  # images of reference frame
                      ref_bboxes=None,  # gt bbox of reference frame
                      gt_pids=None,  # gt ids of current frame bbox mapped to reference frame
                      gt_masks=None,
                      proposals=None,
                      key_frame=None,
                      train_flow=None,
                      key_feature_maps=None,
                      cur_feature_maps=None):
        if train_flow is not None:
            return self.train_flow(img, img_meta, gt_bboxes, gt_bboxes_ignore, gt_labels, ref_img, gt_masks,
                                   key_feature_maps=key_feature_maps, cur_feature_maps=cur_feature_maps)

        return 0, None
        is_first = img_meta[0]['is_first']
        if is_first:
            key_frame = None

        # Extract feature
        x, feat_res0 = self.extract_feat(img)
        ref_x, ref_feat_res0 = self.extract_feat(ref_img)

        losses = dict()

        # RPN forward and loss
        if self.with_rpn:
            rpn_outs = self.rpn_head(x)
            rpn_loss_inputs = rpn_outs + (gt_bboxes, img_meta,
                                          self.train_cfg.rpn)
            rpn_losses = self.rpn_head.loss(*rpn_loss_inputs)
            losses.update(rpn_losses)

            proposal_inputs = rpn_outs + (img_meta, self.test_cfg.rpn)
            proposal_list = self.rpn_head.get_bboxes(*proposal_inputs)
        else:
            proposal_list = proposals

        # assign gts and sample proposals
        if self.with_bbox or self.with_mask:
            bbox_assigner = build_assigner(self.train_cfg.rcnn.assigner)
            bbox_sampler = build_sampler(
                self.train_cfg.rcnn.sampler, context=self)
            num_imgs = img.size(0)
            sampling_results = []
            for i in range(num_imgs):
                assign_result = bbox_assigner.assign(
                    proposal_list[i], gt_bboxes[i], gt_bboxes_ignore[i],
                    gt_labels[i], gt_pids[i])
                sampling_result = bbox_sampler.sample(
                    assign_result,
                    proposal_list[i],
                    gt_bboxes[i],
                    gt_labels[i],
                    gt_pids[i],
                    feats=[lvl_feat[i][None] for lvl_feat in x])
                sampling_results.append(sampling_result)

        # bbox head forward and loss
        if self.with_bbox:
            rois = bbox2roi([res.bboxes for res in sampling_results])
            bbox_img_n = [res.bboxes.size(0) for res in sampling_results]
            ref_rois = bbox2roi(ref_bboxes)
            ref_bbox_img_n = [x.size(0) for x in ref_bboxes]
            # TODO: a more flexible way to decide which feature maps to use
            bbox_feats = self.bbox_roi_extractor(
                x[:self.bbox_roi_extractor.num_inputs], rois)
            ref_bbox_feats = self.bbox_roi_extractor(
                ref_x[:self.bbox_roi_extractor.num_inputs], ref_rois)
            cls_score, bbox_pred = self.bbox_head(bbox_feats)
            # fetch bbox and object_id targets
            bbox_targets, (ids, id_weights) = self.bbox_head.get_target(
                sampling_results, gt_bboxes, gt_labels, self.train_cfg.rcnn)
            loss_bbox = self.bbox_head.loss(cls_score, bbox_pred,
                                            *bbox_targets)
            losses.update(loss_bbox)
            match_score = self.track_head(bbox_feats, ref_bbox_feats,
                                          bbox_img_n, ref_bbox_img_n)
            loss_match = self.track_head.loss(match_score,
                                              ids, id_weights)
            losses.update(loss_match)

        # mask head forward and loss
        if self.with_mask:
            pos_rois = bbox2roi([res.pos_bboxes for res in sampling_results])
            mask_feats = self.mask_roi_extractor(
                x[:self.mask_roi_extractor.num_inputs], pos_rois)
            mask_pred = self.mask_head(mask_feats)

            mask_targets = self.mask_head.get_target(
                sampling_results, gt_masks, self.train_cfg.rcnn)
            pos_labels = torch.cat(
                [res.pos_gt_labels for res in sampling_results])
            loss_mask = self.mask_head.loss(mask_pred, mask_targets,
                                            pos_labels)
            losses.update(loss_mask)

        return losses, x

    def simple_test_bboxes(self,
                           x,
                           img_meta,
                           proposals,
                           rcnn_test_cfg,
                           rescale=False):
        """Test only det bboxes without augmentation."""
        rois = bbox2roi(proposals)
        roi_feats = self.bbox_roi_extractor(
            x[:len(self.bbox_roi_extractor.featmap_strides)], rois)
        cls_score, bbox_pred = self.bbox_head(roi_feats)
        img_shape = img_meta[0]['img_shape']
        scale_factor = img_meta[0]['scale_factor']
        is_first = img_meta[0]['is_first']
        det_bboxes, det_labels = self.bbox_head.get_det_bboxes(
            rois,
            cls_score,
            bbox_pred,
            img_shape,
            scale_factor,
            rescale=rescale,
            cfg=rcnn_test_cfg)
        if det_bboxes.nelement() == 0:
            det_obj_ids = np.array([], dtype=np.int64)
            if is_first:
                self.prev_bboxes = None
                self.prev_roi_feats = None
                self.prev_det_labels = None
            return det_bboxes, det_labels, det_obj_ids

        res_det_bboxes = det_bboxes.clone()
        if rescale:
            res_det_bboxes[:, :4] *= scale_factor

        det_rois = bbox2roi([res_det_bboxes])
        det_roi_feats = self.bbox_roi_extractor(
            x[:self.bbox_roi_extractor.num_inputs], det_rois)
        # recompute bbox match feature

        if is_first or (not is_first and self.prev_bboxes is None):
            det_obj_ids = np.arange(det_bboxes.size(0))
            # save bbox and features for later matching
            self.prev_bboxes = det_bboxes
            self.prev_roi_feats = det_roi_feats
            self.prev_det_labels = det_labels
        else:

            assert self.prev_roi_feats is not None
            # only support one image at a time
            bbox_img_n = [det_bboxes.size(0)]
            prev_bbox_img_n = [self.prev_roi_feats.size(0)]
            match_score = self.track_head(det_roi_feats, self.prev_roi_feats,
                                          bbox_img_n, prev_bbox_img_n)[0]
            match_logprob = torch.nn.functional.log_softmax(match_score, dim=1)
            label_delta = (self.prev_det_labels == det_labels.view(-1, 1)).float()
            bbox_ious = bbox_overlaps(det_bboxes[:, :4], self.prev_bboxes[:, :4])
            # compute comprehensive score
            comp_scores = self.track_head.compute_comp_scores(match_logprob,
                                                              det_bboxes[:, 4].view(-1, 1),
                                                              bbox_ious,
                                                              label_delta,
                                                              add_bbox_dummy=True)
            match_likelihood, match_ids = torch.max(comp_scores, dim=1)
            # translate match_ids to det_obj_ids, assign new id to new objects
            # update tracking features/bboxes of exisiting object,
            # add tracking features/bboxes of new object
            match_ids = match_ids.cpu().numpy().astype(np.int32)
            det_obj_ids = np.ones((match_ids.shape[0]), dtype=np.int32) * (-1)
            best_match_scores = np.ones((self.prev_bboxes.size(0))) * (-100)
            for idx, match_id in enumerate(match_ids):
                if match_id == 0:
                    # add new object
                    det_obj_ids[idx] = self.prev_roi_feats.size(0)
                    self.prev_roi_feats = torch.cat((self.prev_roi_feats, det_roi_feats[idx][None]), dim=0)
                    self.prev_bboxes = torch.cat((self.prev_bboxes, det_bboxes[idx][None]), dim=0)
                    self.prev_det_labels = torch.cat((self.prev_det_labels, det_labels[idx][None]), dim=0)
                else:
                    # multiple candidate might match with previous object, here we choose the one with
                    # largest comprehensive score
                    obj_id = match_id - 1
                    match_score = comp_scores[idx, match_id]
                    if match_score > best_match_scores[obj_id]:
                        det_obj_ids[idx] = obj_id
                        best_match_scores[obj_id] = match_score
                        # udpate feature
                        self.prev_roi_feats[obj_id] = det_roi_feats[idx]
                        self.prev_bboxes[obj_id] = det_bboxes[idx]

        return det_bboxes, det_labels, det_obj_ids

    def full_test(self, img, img_meta, proposals=None, rescale=False):
        """Test without Flow Head."""
        x, _ = self.extract_feat(img)
        key_feat_maps = x

        proposal_list = self.simple_test_rpn(
            x, img_meta, self.test_cfg.rpn) if proposals is None else proposals

        det_bboxes, det_labels, det_obj_ids = self.simple_test_bboxes(
            x, img_meta, proposal_list, self.test_cfg.rcnn, rescale=rescale)
        bbox_results = bbox2result_with_id(det_bboxes, det_labels, det_obj_ids,
                                   self.bbox_head.num_classes)

        if not self.with_mask:
            return bbox_results
        else:
            segm_results = self.simple_test_mask(
                x, img_meta, det_bboxes, det_labels,
                rescale=rescale, det_obj_ids=det_obj_ids)

            return bbox_results, segm_results, key_feat_maps

    def low_test(self, img, img_meta, key_img, key_feats, proposals=None, rescale=False):
        """Test with Flow Head."""
        full_img = key_img
        current_img = img
        feat_last = key_feats
        size_full = full_img.shape
        size_current = current_img.shape
        if size_full[-1] >= size_current[-1]:
            size = size_current[-2:]
            full_img = self.resize(full_img, size)
        else:
            size = size_full[-2:]
            current_img = self.resize(current_img, size)

        # flow
        flow_output = self.flow_head(full_img, current_img)
        upsampling = None
        if upsampling is not None:
            flow_output = F.interpolate(flow_output, size=current_img.size()[-2:], mode='bilinear', align_corners=False)
        # extract feat
        x, _ = self.extract_feat(img)
        # img_size = img.shape[-2:]
        # x = self.warp_feat(feat_last, flow_output*20, (img_size[0], img_size[1]), mode='warp')
        proposal_list = self.simple_test_rpn(
            x, img_meta, self.test_cfg.rpn) if proposals is None else proposals

        det_bboxes, det_labels, det_obj_ids = self.simple_test_bboxes(
            x, img_meta, proposal_list, self.test_cfg.rcnn, rescale=rescale)
        bbox_results = bbox2result_with_id(det_bboxes, det_labels, det_obj_ids,
                                   self.bbox_head.num_classes)

        def feat_fusion(x, feat_last):
            feat = []
            for i in range(len(x)):
                _size = x[i].shape[-2:]
                _tmp = torch.nn.functional.interpolate(feat_last[i], _size, mode='bilinear', align_corners=True)
                # feat.append(torch.max(x[i], _tmp))
                # feat.append((x[i] + _tmp) / 2)
                feat.append((x[i] + _tmp))

            return tuple(feat)
        img_size = img.shape[-2:]
        x = self.warp_feat(feat_last, flow_output*20, (img_size[0], img_size[1]), mode='warp')
        # x_flow = self.warp_feat(feat_last, flow_outputs1)
        # x1 = [resize(x_f, x11.shape[-2:]) for x_f, x11 in zip(x_flow, x)]
        # x1 = feat_fusion(x, x_flow)
        if not self.with_mask:
            return bbox_results
        else:
            segm_results = self.simple_test_mask(
                x, img_meta, det_bboxes, det_labels,
                rescale=rescale, det_obj_ids=det_obj_ids)

            return bbox_results, segm_results, None

    def simple_test(self, img, img_meta, proposals=None, rescale=False, key_frame=None):
        """Test without augmentation."""
        assert self.with_bbox, "Bbox head must be implemented."
        assert self.with_track, "Track head must be implemented"

        is_first = img_meta[0]['is_first']

        # if rl is None:
        if is_first:
            key_frame = None
        if key_frame is not None:
            result = self.low_test(img, img_meta, key_frame['img'], key_frame['feat_map_last'],
                                   proposals=proposals, rescale=rescale)
        else:
            result = self.full_test(img, img_meta, proposals=proposals, rescale=rescale)

        return result

    def aug_test(self, imgs, img_metas, rescale=False):
        """Test with augmentations.

        If rescale is False, then returned bboxes and masks will fit the scale
        of imgs[0].
        """
        # recompute feats to save memory
        proposal_list = self.aug_test_rpn(
            self.extract_feats(imgs), img_metas, self.test_cfg.rpn)
        det_bboxes, det_labels = self.aug_test_bboxes(
            self.extract_feats(imgs), img_metas, proposal_list,
            self.test_cfg.rcnn)

        if rescale:
            _det_bboxes = det_bboxes
        else:
            _det_bboxes = det_bboxes.clone()
            _det_bboxes[:, :4] *= img_metas[0][0]['scale_factor']
        bbox_results = bbox2result(_det_bboxes, det_labels,
                                   self.bbox_head.num_classes)

        # det_bboxes always keep the original scale
        if self.with_mask:
            segm_results = self.aug_test_mask(
                self.extract_feats(imgs), img_metas, det_bboxes, det_labels)
            return bbox_results, segm_results
        else:
            return bbox_results

    def train_flow(self,
                      img,  # images of key frame
                      img_meta,
                      gt_bboxes,
                      gt_bboxes_ignore,
                      gt_labels,
                      cur_img,  # images of current frame
                      gt_masks=None,
                      proposals=None,
                      rescale=False,
                      upsampling = None,
                      key_feature_maps=None,
                      cur_feature_maps=None):
        with torch.autograd.set_detect_anomaly(True):
            losses = dict()
            # extract feature maps of key frame
            if key_feature_maps is not None:
                x_key = key_feature_maps
            else:
                x_key, _ = self.extract_feat(img)

            img = self.resize(img, cur_img.shape[-2:])
            if cur_feature_maps is not None:
                x = cur_feature_maps
            else:
                # extract feature maps of current frame
                x, _ = self.extract_feat(cur_img)

            # detect
            proposal_list = self.simple_test_rpn(
                x, img_meta, self.test_cfg.rpn) if proposals is None else proposals
            rois = bbox2roi(proposal_list, stack=False)
            roi_feats = [self.bbox_roi_extractor(
                x[:len(self.bbox_roi_extractor.featmap_strides)], roi) for roi in rois]
            _results = [self.bbox_head(roi_feat) for roi_feat in roi_feats]
            img_shape = img_meta[0]['img_shape']
            scale_factor = img_meta[0]['scale_factor']
            det_results = [self.bbox_head.get_det_bboxes(
                roi,
                cls_score,
                bbox_pred,
                img_shape,
                scale_factor,
                rescale=rescale,
                cfg=self.test_cfg.rcnn) for (roi, (cls_score, bbox_pred)) in zip(rois, _results)]
            det_bboxes = [det_result[0] for det_result in det_results]
            det_labels = [det_result[1] for det_result in det_results]

            device = det_bboxes[0].device
            for i, det_bbox in enumerate(det_bboxes):
                if len(det_bbox) == 0:
                    det_bboxes[i] = torch.tensor([[0, 0, img_shape[1], img_shape[0], 1]]).to(device).float()
                    det_labels[i] = torch.tensor([[-1]]).to(device).float()

            # calculate flow output
            flow_output = self.flow_head(img, cur_img)
            if upsampling is not None:
                flow_output = F.interpolate(flow_output, size=cur_img.size()[-2:], mode='bilinear', align_corners=False)

            # flow_output.mean().backward()
            # feature warping
            img_size = cur_img.shape[-2:]
            x = self.warp_feat(x_key, flow_output*20, (img_size[0], img_size[1]), mode='warp')

            # mask prediction
            # if det_bboxes is rescaled to the original image size, we need to
            # rescale it back to the testing scale to obtain RoIs.
            scale_factor = img_meta[0]['scale_factor']
            _bboxes = (det_bboxes[:, :4] * scale_factor
                       if rescale else det_bboxes)
            mask_rois = bbox2roi(_bboxes, stack=False)
            mask_feats = [self.mask_roi_extractor(
                x[:len(self.mask_roi_extractor.featmap_strides)], mask_roi) for mask_roi in mask_rois]
            mask_pred = [self.mask_head(mask_feat) for mask_feat in mask_feats if mask_feat.size()[0] > 0]
            # mask_pred[0].backward()

            # mask ground truths
            def assign_gt_single(det_bbox, gt_bbox, pos_iou_thr = 0.5):
                bboxes = det_bbox[:, :4]
                overlaps = bbox_overlaps(gt_bbox, bboxes)
                # for each gt, which predict best overlaps with it
                # for each gt, the max iou of all predictions
                max_overlaps, argmax_overlaps = overlaps.max(dim=1)
                num_gts, num_bboxes = overlaps.size(0), overlaps.size(1)
                assigned_gt_inds = overlaps.new_full((num_gts,), -1, dtype=torch.long)
                # assign positive: above positive IoU threshold
                pos_inds = max_overlaps >= pos_iou_thr
                assigned_gt_inds[pos_inds] = argmax_overlaps[pos_inds]
                return assigned_gt_inds

            assigned_gt_inds = [assign_gt_single(det_bbox, gt_bbox) for (det_bbox, gt_bbox) in zip(det_bboxes, gt_bboxes)]

            def mask_target_single(gt_bboxes, gt_masks, cfg):
                # get mask(28X28) for each bbox.
                mask_size = cfg.mask_size
                num_gt = gt_bboxes.size(0)
                mask_targets = []
                if num_gt > 0:
                    gt_bboxes_np = gt_bboxes.cpu().numpy()
                    # pos_assigned_gt_inds = pos_assigned_gt_inds.cpu().numpy()
                    for i in range(num_gt):
                        gt_mask = gt_masks[i]
                        bbox = gt_bboxes_np[i, :].astype(np.int32)
                        x1, y1, x2, y2 = bbox
                        w = np.maximum(x2 - x1 + 1, 1)
                        h = np.maximum(y2 - y1 + 1, 1)
                        # mask is uint8 both before and after resizing
                        target = mmcv.imresize(gt_mask[y1:y1 + h, x1:x1 + w],
                                               (mask_size, mask_size))
                        mask_targets.append(target)
                    mask_targets = torch.from_numpy(np.stack(mask_targets)).float().to(
                        gt_bboxes.device)
                else:
                    mask_targets = gt_bboxes.new_zeros((0, mask_size, mask_size))
                return mask_targets

            mask_targets = torch.cat(
                [mask_target_single(gt_bbox, gt_mask, self.train_cfg.rcnn) for (gt_bbox, gt_mask) in zip(gt_bboxes, gt_masks)])

            mask_size = self.train_cfg.rcnn.mask_size
            mask_p = []
            pos_labels = []

            # assign predicted mask and label to each mask target.
            for i in range(len(mask_pred)):
                assigned_gt_ind = assigned_gt_inds[i]
                for id in assigned_gt_ind:
                    if id == -1:
                        mask = mask_pred[i].new_zeros((1, 41, mask_size, mask_size))
                        # label = torch.tensor(0).long().to(device)
                        label = gt_labels[0].new_zeros((1,))
                    else:
                        mask = mask_pred[i][id].unsqueeze(0)
                        label = det_labels[i][id] + 1
                    mask_p.append(mask)
                    pos_labels.append(torch.tensor([label]).long().to(device))
            mask_p = torch.cat(mask_p)
            # pos_labels = torch.cat(pos_labels)
            # TODO pos_labels = gt_labels or det_labels?
            pos_labels = torch.cat(gt_labels)

            # calculate loss
            loss_mask = self.mask_head.loss(mask_p, mask_targets,
                                            pos_labels)
            losses.update(loss_mask)
        torch.cuda.empty_cache()
        return losses

    @staticmethod
    def resize(feat_map, size=(48, 64)):
        """Resize feature map to certain size."""
        key_feature = torch.nn.functional.interpolate(feat_map, size, mode='bilinear', align_corners=False)
        return key_feature

    def warp_feat(self, feat_maps, flow, img_size, feat_stride=(4, 8, 16, 32, 64), mode='warp'):
        assert mode in ['warp', 'gather'], "Ensure mode in ['warp', 'gather']."
        if len(flow.shape) != 4:
            flow = flow.unsqueeze(0)
        sizes = [(int(img_size[0]/stride), int(img_size[1]/stride)) for stride in feat_stride]
        if feat_maps[0].shape[-2:] != sizes[0]:
            feat_maps = [self.resize(feat, size) if feat.shape[-2:] != size else feat for (feat, size) in zip(feat_maps, sizes)]
        if mode == 'warp':
            x = [warp(feat_map, flow / stride) for (feat_map, stride) in zip(feat_maps, feat_stride)]
        elif mode == 'gather':
            from mmdet.models.flow_heads.propagation_utils_deprecated import gather_nd
            flows = [self.resize(flow, size) for (size) in sizes]
            x = [gather_nd(feat_map, flow / stride) for (feat_map, flow, stride) in zip(feat_maps, flows, feat_stride)]

        return x

max_stars_repo_name: ruyueshuo/MaskTrackRCNN

max_issues_repo_name: ruyueshuo/MaskTrackRCNN

max_forks_repo_name: ruyueshuo/MaskTrackRCNN


text:
import unittest
from datetime import datetime

import numpy as np
from dateutil.tz import tzlocal
import ipywidgets as widgets
from nwbwidgets.utils.units import (
    get_min_spike_time,
    align_by_trials,
    align_by_time_intervals,
)
from pynwb import NWBFile
from pynwb.epoch import TimeIntervals

from ..base import TimeIntervalsSelector
from ..misc import TuningCurveWidget, TuningCurveExtendedWidget

class UnitsTrialsTestCase(unittest.TestCase):

    def setUp(self):
        start_time = datetime(2017, 4, 3, 11, tzinfo=tzlocal())
        create_date = datetime(2017, 4, 15, 12, tzinfo=tzlocal())

        self.nwbfile = NWBFile(
            session_description="NWBFile for PSTH",
            identifier="NWB123",
            session_start_time=start_time,
            file_create_date=create_date,
        )

        self.nwbfile.add_unit_column("location", "the anatomical location of this unit")
        self.nwbfile.add_unit_column(
            "quality", "the quality for the inference of this unit"
        )

        self.nwbfile.add_unit(
            id=1,
            spike_times=[2.2, 3.0, 4.5],
            obs_intervals=[[1, 10]],
            location="CA1",
            quality=0.95,
        )
        self.nwbfile.add_unit(
            id=2,
            spike_times=[2.2, 3.0, 25.0, 26.0],
            obs_intervals=[[1, 10], [20, 30]],
            location="CA3",
            quality=0.85,
        )
        self.nwbfile.add_unit(
            id=3,
            spike_times=[1.2, 2.3, 3.3, 4.5],
            obs_intervals=[[1, 10], [20, 30]],
            location="CA1",
            quality=0.90,
        )

        self.nwbfile.add_trial_column(
            name="stim", description="the visual stimuli during the trial"
        )

        self.nwbfile.add_trial(start_time=0.0, stop_time=2.0, stim="person")
        self.nwbfile.add_trial(start_time=3.0, stop_time=5.0, stim="ocean")
        self.nwbfile.add_trial(start_time=6.0, stop_time=8.0, stim="desert")
        self.nwbfile.add_trial(start_time=8.0, stop_time=12.0, stim="person")
        self.nwbfile.add_trial(start_time=13.0, stop_time=15.0, stim="ocean")
        self.nwbfile.add_trial(start_time=16.0, stop_time=18.0, stim="desert")

class ExtendedTimeIntervalSelector(TimeIntervalsSelector):
    InnerWidget = TuningCurveWidget

class ExtendedTimeIntervalSelectorTestCase(UnitsTrialsTestCase):

    def setUp(self):
        super().setUp()

        # add intervals to nwbfile
        ti1 = TimeIntervals(name='intervals', description='experimental intervals')
        ti1.add_interval(start_time=0.0, stop_time=2.0)
        ti1.add_interval(start_time=2.0, stop_time=4.0)
        ti1.add_interval(start_time=4.0, stop_time=6.0)
        ti1.add_interval(start_time=6.0, stop_time=8.0)
        ti1.add_column(name='var1', data=['a', 'b', 'a', 'b'], description='no description')
        self.nwbfile.add_time_intervals(ti1)

        self.widget = ExtendedTimeIntervalSelector(
            input_data=self.nwbfile.units
        )

    def test_make_widget(self):
        assert isinstance(self.widget, widgets.Widget)

    def test_widget_children(self):
        assert len(self.widget.children) == 2

        for i, c in enumerate(self.widget.children):
            assert isinstance(c, widgets.Widget), f'{i}th child of TuningCurve widget is not a widget'

    def test_make_graph(self):
        # rows controller triggers drawing of graphic
        self.widget.children[1].children[1].value = 'var1'

        for i, c in enumerate(self.widget.children):
            assert isinstance(c, widgets.Widget), f'{i}th child of TuningCurve widget is not a widget'

class TuningCurveTestCase(UnitsTrialsTestCase):

    def setUp(self):
        super().setUp()
        self.widget = TuningCurveWidget(
            units=self.nwbfile.units,
            trials=self.nwbfile.trials
        )
        # rows controller triggers drawing of graphic
        self.widget.children[0].children[1].value = 'stim'

    def test_make_widget(self):
        assert isinstance(self.widget, widgets.Widget)

    def test_widget_children(self):
        assert len(self.widget.children) == 2

        for i, c in enumerate(self.widget.children):
            assert isinstance(c, widgets.Widget), f'{i}th child of TuningCurve widget is not a widget'

class TuningCurveRasterGridCombinedTestCase(UnitsTrialsTestCase):

    def setUp(self):
        super().setUp()
        self.widget = TuningCurveExtendedWidget(
            units=self.nwbfile.units,
            trials=self.nwbfile.trials
        )

    def test_make_widget(self):
        assert isinstance(self.widget, widgets.Widget)

    def test_widget_children(self):
        assert len(self.widget.children) == 3

        for i, c in enumerate(self.widget.children):
            assert isinstance(c, widgets.Widget), f'{i}th child of TuningCurve widget is not a widget'

class ShowPSTHTestCase(UnitsTrialsTestCase):

    def test_get_min_spike_time(self):
        assert get_min_spike_time(self.nwbfile.units) == 1.2

    def test_align_by_trials(self):
        compare_to_at = [
            np.array([2.2, 3.0, 25.0, 26.0]),
            np.array([-0.8, 0.0, 22.0, 23.0]),
            np.array([-3.8, -3.0, 19.0, 20.0]),
            np.array([-5.8, -5., 17., 18.]),
            np.array([-10.8, -10.,  12.,  13.]),
            np.array([-13.8, -13.,   9.,  10.])
        ]

        at = align_by_trials(
            self.nwbfile.units,
            index=1,
            start=-20.0,
            end=30.0
        )

        np.testing.assert_allclose(at, compare_to_at, rtol=1e-02)

    def test_align_by_time_intervals_Nonetrials_select(self):
        time_intervals = TimeIntervals(name="Test Time Interval")
        time_intervals.add_interval(start_time=21.0, stop_time=28.0)
        time_intervals.add_interval(start_time=22.0, stop_time=26.0)
        time_intervals.add_interval(start_time=22.0, stop_time=28.4)

        ati = align_by_time_intervals(
            self.nwbfile.units,
            index=1,
            intervals=time_intervals,
            stop_label=None,
            start=-20.0,
            end=30.0,
        )

        compare_to_ati = [
            np.array([-18.8, -18.0, 4.0, 5.0]),
            np.array([-19.8, -19.0, 3.0, 4.0]),
            np.array([-19.8, -19.0, 3.0, 4.0]),
        ]

        np.testing.assert_array_equal(ati, compare_to_ati)

    def test_align_by_time_intervals(self):
        time_intervals = TimeIntervals(name="Test Time Interval")
        time_intervals.add_interval(start_time=21.0, stop_time=28.0)
        time_intervals.add_interval(start_time=22.0, stop_time=26.0)
        time_intervals.add_interval(start_time=22.0, stop_time=28.4)

        ati = align_by_time_intervals(
            self.nwbfile.units,
            index=1,
            intervals=time_intervals,
            stop_label=None,
            start=-20.0,
            end=30.0,
            rows_select=[0, 1],
        )

        compare_to_ati = [
            np.array([-18.8, -18.0, 4.0, 5.0]),
            np.array([-19.8, -19.0, 3.0, 4.0]),
        ]

        np.testing.assert_array_equal(ati, compare_to_ati)

ext: py

max_issues_repo_name: catalystneuro/nwb-jupyter-widgets

max_forks_repo_name: catalystneuro/nwb-jupyter-widgets

reason: import numpy


text:
import os
import random
from tempfile import NamedTemporaryFile, TemporaryDirectory
from typing import Dict, List

import numpy as np
import torch

from lhotse import (
    AudioSource,
    CutSet,
    Fbank,
    FbankConfig,
    LilcomFilesWriter,
    LilcomHdf5Writer,
    MonoCut,
    NumpyHdf5Writer,
    Recording,
    SupervisionSegment,
)
from lhotse.array import seconds_to_frames
from lhotse.supervision import AlignmentItem
from lhotse.utils import Seconds, uuid4

def random_cut_set(n_cuts=100) -> CutSet:
    return CutSet.from_cuts(
        MonoCut(
            id=uuid4(),
            start=round(random.uniform(0, 5), ndigits=8),
            duration=round(random.uniform(3, 10), ndigits=8),
            channel=0,
            recording=Recording(
                id=uuid4(),
                sources=[],
                sampling_rate=16000,
                num_samples=1600000,
                duration=100.0,
            ),
        )
        for _ in range(n_cuts)
    )

class RandomCutTestCase:
    def setup_method(self, method):
        self.files = []
        self.dirs = []

    def teardown_method(self, method):
        self.cleanup()

    def cleanup(self):
        for f in self.files:
            f.close()
        self.files = []
        for d in self.dirs:
            d.cleanup()
        self.dirs = []

    def with_recording(
        self, sampling_rate: int, num_samples: int, use_zeros: bool = False
    ) -> Recording:
        import torchaudio  # torchaudio does not have issues on M1 macs unlike soundfile

        f = NamedTemporaryFile("wb", suffix=".wav")
        self.files.append(f)
        duration = num_samples / sampling_rate
        if use_zeros:
            samples = torch.zeros((1, num_samples))
        else:
            samples = torch.rand((1, num_samples))
        torchaudio.save(f.name, samples, sample_rate=sampling_rate)
        f.flush()
        os.fsync(f)
        return Recording(
            id=str(uuid4()),
            sources=[AudioSource(type="file", channels=[0], source=f.name)],
            sampling_rate=sampling_rate,
            num_samples=num_samples,
            duration=duration,
        )

    def with_cut(
        self,
        sampling_rate: int,
        num_samples: int,
        features: bool = True,
        supervision: bool = False,
        alignment: bool = False,
        custom_field: bool = False,
        frame_shift: Seconds = 0.01,
        use_zeroes: bool = False,
    ) -> MonoCut:
        duration = num_samples / sampling_rate
        cut = MonoCut(
            id=str(uuid4()),
            start=0,
            duration=duration,
            channel=0,
            recording=self.with_recording(
                sampling_rate=sampling_rate,
                num_samples=num_samples,
                use_zeros=use_zeroes,
            ),
        )
        if features:
            cut = self._with_features(
                cut, frame_shift=frame_shift, sampling_rate=sampling_rate
            )
        if supervision:
            cut.supervisions.append(
                SupervisionSegment(
                    id=f"sup-{cut.id}",
                    recording_id=cut.recording_id,
                    start=0,
                    duration=cut.duration,
                    text="irrelevant",
                    alignment=self._with_alignment(cut, "irrelevant")
                    if alignment
                    else None,
                )
            )
        if custom_field:
            self._with_custom_temporal_array(cut=cut, frame_shift=frame_shift)
        return cut

    def _with_features(
        self, cut: MonoCut, frame_shift: Seconds, sampling_rate: int
    ) -> MonoCut:
        d = TemporaryDirectory()
        self.dirs.append(d)
        extractor = Fbank(
            config=FbankConfig(sampling_rate=sampling_rate, frame_shift=frame_shift)
        )
        with LilcomHdf5Writer(d.name) as storage:
            return cut.compute_and_store_features(extractor, storage=storage)

    def _with_alignment(
        self, cut: MonoCut, text: str
    ) -> Dict[str, List[AlignmentItem]]:
        subwords = [
            text[i : i + 3] for i in range(0, len(text), 3)
        ]  # Create subwords of 3 chars
        dur = cut.duration / len(subwords)
        alignment = [
            AlignmentItem(symbol=sub, start=i * dur, duration=dur)
            for i, sub in enumerate(subwords)
        ]
        return {"subword": alignment}

    def _with_custom_temporal_array(self, cut: MonoCut, frame_shift: Seconds) -> None:
        d = TemporaryDirectory()
        self.dirs.append(d)
        num_frames = seconds_to_frames(cut.duration, frame_shift=frame_shift)
        array = np.random.randint(256, size=(num_frames,))
        with NumpyHdf5Writer(d.name) as storage:
            cut.codebook_indices = storage.store_array(
                key="ali1", value=array, frame_shift=frame_shift, temporal_dim=0
            )

max_stars_repo_name: glynpu/lhotse

max_issues_repo_name: glynpu/lhotse

max_forks_repo_name: glynpu/lhotse


text:
"""
clusters module for electrification-planner

Provides functions to read in a raster population dataset
and convert to discrete vector polgons, each with a set
population value. Additionally calculate each polygon's
distance from a provided grid infrastructure vector.
"""

import json
from pathlib import Path

import numpy as np
from scipy import ndimage
import geopandas as gpd

import rasterio
from rasterio.mask import mask
from rasterio.features import shapes, rasterize
from rasterstats import zonal_stats

def clip_raster(raster, boundary, boundary_layer='gadm36_UGA_0'):
	"""
	Clip the raster to the given administrative boundary.

	Parameters
	----------
	raster: string, pathlib.Path or rasterio.io.DataSetReader
		Location of or already opened raster.
	boundary: string, pathlib.Path or geopandas.GeoDataFrame
		The poylgon by which to clip the raster.
	boundary_layer: string, optional
		For multi-layer files (like GeoPackage), specify the layer to be used.

	Returns
	-------
	tuple
		Three elements:
			clipped: numpy.ndarray
				Contents of clipped raster.
			affine: affine.Affine()
				Information for mapping pixel coordinates
				to a coordinate system.
			crs: dict
				Dict of the form {'init': 'epsg:4326'} defining the coordinate
				reference system of the raster.

	"""

	if isinstance(raster, Path):
		raster = str(raster)
	if isinstance(raster, str):
		raster = rasterio.open(raster)

	crs = raster.crs

	if isinstance(boundary, Path):
		boundary = str(boundary)
	if isinstance(boundary, str):
		if '.gpkg' in boundary:
			driver = 'GPKG'
		else:
			driver = None  # default to shapefile
			boundary_layer = ''  # because shapefiles have no layers

		boundary = gpd.read_file(boundary, layer=boundary_layer, driver=driver)

	boundary = boundary.to_crs(crs=raster.crs)
	coords = [json.loads(boundary.to_json())['features'][0]['geometry']]

	# mask/clip the raster using rasterio.mask
	clipped, affine = mask(dataset=raster, shapes=coords, crop=True)

	return clipped, affine, crs

def create_clusters(raster, affine, crs):
	"""
	Create a polygon GeoDataFrame from the given raster

	Parameters
	----------
	raster: numpy.ndarray
		The raster data to use.
	affine: affine.Affine()
		Raster pixel mapping information.
	crs: dict
		Dict of the form {'init': 'epsg:4326'} defining the coordinate
		reference system to use.

	Returns
	-------
	clusters: geopandas.GeoDataFrame
		A GeoDataFrame with integer index and two columns:
		geometry contains the Shapely polygon representations
		raster_val contains the values from the raster

	"""

	geoms = list(({'properties': {'raster_val': v}, 'geometry': s}
	              for i, (s, v)
	              in enumerate(shapes(raster, mask=None, transform=affine))))

	clusters = gpd.GeoDataFrame.from_features(geoms)
	clusters.crs = crs

	return clusters

# TODO Could instead filter at the raster stage?
def filter_merge_clusters(clusters, max_block_size_multi=5, min_block_pop=50, buffer_amount=150):
	"""
	The vectors created by create_clusters() are a single square for each raster pixel.
	This function does the follows:
	- Remove overly large clusters, caused by defects in the input raster.
	- Remove clusters with population below a certain threshold.
	- Buffer the remaining clusters and merge those that overlap.

	Parameters
	----------
	clusters: geopandas.GeoDataFrame
		The unprocessed clusters created by create_clusters()
	max_block_size_multi: int, optional
		Remove clusters that are more than this many times average size. Default 5.
	min_block_pop: int, optional
		Remove clusters with below this population. Default 50.
	buffer_amount: int, optional
		Distance in metres by which to buffer the clusters before merging. Default 150.

	Returns
	-------
	clusters: geopandas.GeoDataFrame
		The processed clusters.
	"""

	# remove blocks that are too big (basically artifacts)
	clusters['area_m2'] = clusters.geometry.area
	clusters = clusters[clusters['area_m2'] < clusters['area_m2'].mean() * max_block_size_multi]

	# remove blocks with too few people
	clusters = clusters[clusters['raster_val'] > min_block_pop]

	# buffer outwards so that nearby blocks will overlap
	clusters['geometry'] = clusters.geometry.buffer(buffer_amount)

	# and dissolve the thousands of blocks into a single layer (with no attributes!)
	clusters['same'] = 1
	clusters = clusters.dissolve(by='same')

	# To get our attributes back, we convert the dissolves polygon into singleparts
	# This means each contiguous bubble becomes its own polygon and can store its own attributes
	crs = clusters.crs
	clusters = clusters.explode()
	clusters = clusters.reset_index()

	# no longer needed in GeoPandas >= 0.4.0
	# clusters['geometry'] = clusters[0]
	# clusters = gpd.GeoDataFrame(clusters)
	# clusters.crs = crs

	clusters = clusters.drop(columns=['same', 'level_1', 'raster_val'])  # raster_val is no longer meaningful

	# And then add the polygon's area back to its attributes
	clusters["area_m2"] = clusters['geometry'].area

	return clusters

def cluster_pops(clusters, raster, affine=None):
	"""
	The filter_merge_clusters() process loses the underlying raster values.
	So we need to use rasterstats.zonal_stats() to get it back.

	Parameters
	----------
	clusters: geopandas.GeoDataFrame
		The processed clusters.
	raster: str, pathlib.Path or numpy.ndarray
		Either a path to the raster, or an already imported numpy.ndarray with the data.
	affine: affine.Affine(), optional
		If a numpy ndarray is passed above, the affine is also needed.

	Returns
	-------
	clusters: geopandas.GeoDataFrame
		The processed clusters.
	"""
	if isinstance(raster, Path):
		raster = str(raster)
	if isinstance(raster, str):
		pop_sums = zonal_stats(clusters, raster, stats='sum')

	else:
		pop_sums = zonal_stats(clusters, raster, affine=affine, stats='sum', nodata=0)

	clusters['pop_sum'] = [x['sum'] for x in pop_sums]

	return clusters

def cluster_grid_distance(clusters, grid, shape, affine):
	"""
	Use a vector containing grid infrastructure to determine
	each cluster's distance from the grid.

	Parameters
	----------
	clusters: geopandas.GeoDataFrame
		The processed clusters.
	grid: str, pathlib.Path or geopandas.GeoDataFrame
		Path to or already imported grid dataframe.
	shape: tuple
		Tuple of two integers representing the shape of the data
		for rasterizing grid. Sould match the clipped raster.
	affine: affine.Affine()
		As above, should match the clipped raster.

	Returns
	-------
	clusters: geopandas.GeoDataFrame
		The processed clusters.

	"""

	if isinstance(grid, Path):
		grid = str(grid)
	if isinstance(grid, str):
		grid = gpd.read_file(grid)

	grid = grid.to_crs(crs=clusters.crs)
	grid = grid.loc[grid['geometry'].length > 0]

	grid_raster = rasterize(grid.geometry, out_shape=shape, fill=1,
	                        default_value=0, all_touched=True, transform=affine)
	dist_raster = ndimage.distance_transform_edt(grid_raster) * affine[0]

	dists = zonal_stats(vectors=clusters, raster=dist_raster, affine=affine, stats='min', nodata=1000)
	clusters['grid_dist'] = [x['min'] for x in dists]

	return clusters

def save_clusters(clusters, out_path):
	"""
	Convert to EPSG:4326 and save to the specified file.
	clusters: geopandas.GeoDataFrame
		The processed clusters.
	out_path: str or pathlib.Path
		Where to save the clusters file.
	"""

	if isinstance(out_path, Path):
		out_path = str(out_path)
	if '.gpkg' in out_path:
		driver = 'GPKG'
	else:
		driver = None

	clusters = clusters.to_crs(epsg=4326)
	clusters.to_file(out_path, driver=driver)

max_stars_repo_name: carderne/electrification-planner

max_issues_repo_name: carderne/electrification-planner

max_forks_repo_name: carderne/electrification-planner

reason: import numpy,from scipy


text:
"""
Do volume, salt, and salt-squared budgets for user-specified volumes.
This is a fairly complicated task because it involves coordinating
the net salt and volume time series from flux_[get,lowpass]_s.py,
the TEF transports through any open sections of the volume,
and all the rivers flowing into the volume.

Run with a command like:
run flux_salt_budget
and then it runs for all volumes in vol_list and all years in year_list.

Currently need to do more extractions to support 2018 and 2019.

"""

import os; import sys
sys.path.append(os.path.abspath('../alpha'))
import Lfun
import tef_fun
import flux_fun
from importlib import reload
reload(flux_fun)

import matplotlib.pyplot as plt
import numpy as np
import pickle
import pandas as pd
from datetime import datetime, timedelta
import argparse

# debugging imports
import netCDF4 as nc
from time import time
import zfun

parser = argparse.ArgumentParser()
parser.add_argument('-g', '--gridname', type=str, default='cas6')
parser.add_argument('-t', '--tag', type=str, default='v3')
parser.add_argument('-x', '--ex_name', type=str, default='lo8b')
parser.add_argument('-test', '--testing', type=zfun.boolean_string, default=False)
args = parser.parse_args()
testing = args.testing

old_dt = False

# Get Ldir
Ldir = Lfun.Lstart(args.gridname, args.tag)
gtagex = args.gridname + '_' + args.tag + '_' + args.ex_name

if testing:
    year_list = [2019]
    vol_list = ['Puget Sound']
    save_figs = False
else:
    year_list = [2017, 2018, 2019]
    vol_list = ['Salish Sea', 'Puget Sound', 'Hood Canal']
    save_figs = True

# make DataFrames to hold error statistics
err_df_vol = pd.DataFrame(index=year_list, columns=vol_list)
err_df_salt = pd.DataFrame(index=year_list, columns=vol_list)
err_df_salt2 = pd.DataFrame(index=year_list, columns=vol_list)

plt.close('all')
for which_vol in vol_list:

    for year in year_list:
        year_str = str(year)

        # select input/output location
        run_name = gtagex+'_'+year_str+'.01.01_'+year_str+'.12.31'
        indir00 = Ldir['LOo'] + 'tef2/'
        indir0 = indir00 + run_name + '/'

        outdir = indir00 + 'salt_budget_plots/'
        Lfun.make_dir(outdir)

        # load low passed segment volume, net salt, and other DataFrames
        v_lp_df = pd.read_pickle(indir0 + 'flux/daily_segment_volume.p')
        sv_lp_df = pd.read_pickle(indir0 + 'flux/daily_segment_net_salt.p')
        s2v_lp_df = pd.read_pickle(indir0 + 'flux/daily_segment_net_salt2.p')
        mix_lp_df = pd.read_pickle(indir0 + 'flux/daily_segment_mix.p')
        hmix_lp_df = pd.read_pickle(indir0 + 'flux/daily_segment_hmix.p')

        # new
        vt_lp_df = pd.read_pickle(indir0 + 'flux/daily_segment_vt.p')
        svt_lp_df = pd.read_pickle(indir0 + 'flux/daily_segment_svt.p')
        s2vt_lp_df = pd.read_pickle(indir0 + 'flux/daily_segment_s2vt.p')

        # Info specific to each volume
        # The sign for each section indicates which direction is INTO the volume.
        if which_vol == 'Salish Sea':
            seg_list = list(v_lp_df.columns)
            sect_sign_dict = {'jdf1':1, 'sog5':-1}
        elif which_vol == 'Puget Sound':
            seg_list = (flux_fun.ssA + flux_fun.ssM + flux_fun.ssT
                + flux_fun.ssS + flux_fun.ssW + flux_fun.ssH)
            sect_sign_dict = {'ai1':1, 'dp':1}
        elif which_vol == 'Hood Canal':
            seg_list = flux_fun.ssH
            sect_sign_dict = {'hc1':-1}

        v_lp_df = v_lp_df[seg_list]
        sv_lp_df = sv_lp_df[seg_list]
        s2v_lp_df = s2v_lp_df[seg_list]

        vt_lp_df = vt_lp_df[seg_list]
        svt_lp_df = svt_lp_df[seg_list]
        s2vt_lp_df = s2vt_lp_df[seg_list]

        mix_lp_df = mix_lp_df[seg_list]
        hmix_lp_df = hmix_lp_df[seg_list]

        sect_df = tef_fun.get_sect_df()

        river_list = []
        for seg_name in seg_list:
            seg = flux_fun.segs[seg_name]
            river_list = river_list + seg['R']
        riv_df = pd.read_pickle(Ldir['LOo'] + 'river/'
            + Ldir['gtag'] + '_'+year_str+'.01.01_'+year_str+'.12.31.p')
        riv_df.index += timedelta(days=0.5)
        riv_df = riv_df[river_list]

        tef_df_dict = {}
        for sn in sect_sign_dict.keys():
            sect_sign = sect_sign_dict[sn]
            tef_df_dict[sn], in_sign = flux_fun.get_fluxes(indir0, sn)
            if in_sign != sect_sign:
                print('WARNING: potential sign error!!')

        # BUDGETS
        indall = sv_lp_df.index

        # miscellaneous stuff that doesn't go in budget
        misc_df = pd.DataFrame(0,index=indall, columns = ['Qprism','Ftide'])
        # we initialize only columns where we sum over TEF sections
        for sect_name in tef_df_dict.keys():
            df = tef_df_dict[sect_name]
            misc_df['Qprism'] = misc_df['Qprism'] + df['Qtide']/2 # [m3/s]
            misc_df['Ftide'] = misc_df['Ftide'] + df['Ftide'] # [Watts]
        v = v_lp_df.sum(axis=1).to_numpy()
        vt = vt_lp_df.sum(axis=1).to_numpy()
        svt = svt_lp_df.sum(axis=1).to_numpy()
        s2vt = s2vt_lp_df.sum(axis=1).to_numpy()
        misc_df['V'] = v # [m3]

        # also create, from the hourly files, <s> and d/dt (int <s>2 dV)
        vh_df = pd.read_pickle(indir0 + 'flux/hourly_segment_volume.p')
        sh_df = pd.read_pickle(indir0 + 'flux/hourly_segment_salinity.p')
        vh_df = vh_df[seg_list]
        sh_df = sh_df[seg_list]
        svh_df = sh_df * vh_df # net salt in each segment (hourly)
        smeanh = svh_df.sum(axis=1).to_numpy() / vh_df.sum(axis=1).to_numpy()
        smean2h = smeanh * smeanh
        smean2vh = smean2h * vh_df.sum(axis=1).to_numpy()
        dsmean2vh_dt = np.nan + np.ones_like(smean2vh)
        dsmean2vh_dt[1:-1] = (smean2vh[2:] - smean2vh[:-2])/(2*3600)
        dsmean2vh_dt_lp = zfun.filt_godin(dsmean2vh_dt)
        pad = 36
        dsmean2v_dt_lp = dsmean2vh_dt_lp[pad:-(pad+1):24]

        if old_dt:
            misc_df['Smean'] = sv_lp_df.sum(axis=1)/misc_df['V']
        else:
            smeanh_lp = zfun.filt_godin(smeanh)
            smean_lp = smeanh_lp[pad:-(pad+1):24]
            misc_df['Smean'] = smean_lp

        # volume budget
        vol_df = pd.DataFrame(0, index=indall, columns=['Qin','Qout'])
        for sect_name in tef_df_dict.keys():
            df = tef_df_dict[sect_name]
            vol_df['Qin'] = vol_df['Qin'] + df['Qin']
            vol_df['Qout'] = vol_df['Qout'] + df['Qout']
        vol_df['Qr'] = riv_df.sum(axis=1)
        if old_dt:
            vol_df.loc[1:-1, 'dV_dt'] = (v[2:] - v[:-2]) / (2*86400)
        else:
            vol_df.loc[:, 'dV_dt'] = vt

        vol_df['Error'] = vol_df['dV_dt'] - vol_df.loc[:,'Qin'] - vol_df.loc[:,'Qout'] - vol_df.loc[:,'Qr']
        vol_rel_err = vol_df['Error'].mean()/vol_df['Qr'].mean()

        # salt budget
        salt_df = pd.DataFrame(0, index=indall, columns=['QSin','QSout'])
        for sect_name in tef_df_dict.keys():
            df = tef_df_dict[sect_name]
            salt_df['QSin'] = salt_df['QSin'] + df['QSin']
            salt_df['QSout'] = salt_df['QSout'] + df['QSout']
        sn = sv_lp_df[seg_list].sum(axis=1).values
        if old_dt:
            salt_df.loc[1:-1, 'dSnet_dt'] = (sn[2:] - sn[:-2]) / (2*86400)
        else:
            salt_df.loc[:, 'dSnet_dt'] = svt

        salt_df['Error'] = salt_df['dSnet_dt'] - salt_df['QSin'] - salt_df['QSout']
        salt_rel_err = salt_df['Error'].mean()/salt_df['QSin'].mean()

        misc_df['Sin'] = salt_df['QSin']/vol_df['Qin']
        misc_df['Sout'] = salt_df['QSout']/vol_df['Qout']
        Socn = misc_df['Sin'].max() # used in mixedness budget

        # re-express the salt budget using Qe notation
        salt_qe_df = pd.DataFrame(index=indall)
        salt_qe_df['dSnet_dt'] = salt_df['dSnet_dt']

        misc_df['Qe'] = (vol_df['Qin'] - vol_df['Qout'])/2
        misc_df['Qnet'] = -(vol_df['Qout'] + vol_df['Qin']) # like Qr, but accounting for dV/dt
        misc_df['DS'] = misc_df['Sin'] - misc_df['Sout']
        misc_df['Sbar'] = (misc_df['Sin'] + misc_df['Sout'])/2
        # NOTE: Sbar is bad notation here, need to coordinate with Variance Budget usage of Smean.

        salt_qe_df['QeDS'] = misc_df['Qe']*misc_df['DS']
        salt_qe_df['-QrSbar'] = -misc_df['Qnet']*misc_df['Sbar']
        salt_qe_df['Error'] = salt_qe_df['dSnet_dt'] - salt_qe_df['QeDS'] - salt_qe_df['-QrSbar']
        salt_qe_rel_err_qe = salt_qe_df['Error'].mean()/salt_qe_df['QeDS'].mean()

        # salinity-squared budget
        salt2_df = pd.DataFrame(0, index=indall, columns=['QS2in','QS2out'])
        for sect_name in tef_df_dict.keys():
            df = tef_df_dict[sect_name]
            salt2_df['QS2in'] = salt2_df['QS2in'] + df['QS2in']
            salt2_df['QS2out'] = salt2_df['QS2out'] + df['QS2out']
        s2n = s2v_lp_df[seg_list].sum(axis=1).to_numpy()
        if old_dt:
            ds2net_dt = (s2n[2:] - s2n[:-2]) / (2*86400)
            salt2_df.loc[1:-1, 'dS2net_dt'] = ds2net_dt
        else:
            ds2net_dt = s2vt
            salt2_df.loc[:, 'dS2net_dt'] = ds2net_dt

        salt2_df['VMix'] = mix_lp_df[seg_list].sum(axis=1).to_numpy()
        salt2_df['HMix'] = hmix_lp_df[seg_list].sum(axis=1).to_numpy()
        salt2_df['Mix_resolved'] = salt2_df['VMix'] + salt2_df['HMix']

        misc_df['S2in'] = salt2_df['QS2in']/vol_df['Qin']
        misc_df['S2out'] = salt2_df['QS2out']/vol_df['Qout']
        misc_df['S2mean'] = s2v_lp_df.sum(axis=1)/misc_df['V']

        salt2_df['Mix_numerical'] = salt2_df['dS2net_dt'] - salt2_df['QS2in'] - salt2_df['QS2out'] - salt2_df['Mix_resolved']
        salt2_df['Mixing'] = salt2_df['Mix_numerical'] + salt2_df['Mix_resolved']
        salt2_rel_err = salt2_df['Mix_numerical'].mean()/salt2_df['Mix_resolved'].mean()

        # variance budget for S'^2
        sp2_df = pd.DataFrame(index=indall)

        misc_df['Sp2in'] = misc_df['S2in'] - 2*misc_df['Smean']*misc_df['Sin'] + misc_df['Smean']**2
        misc_df['Sp2out'] = misc_df['S2out'] - 2*misc_df['Smean']*misc_df['Sout'] + misc_df['Smean']**2

        sp2_df['QSp2in'] = misc_df['Sp2in'] * vol_df['Qin']
        sp2_df['QSp2out'] = misc_df['Sp2out'] * vol_df['Qout']
        sp2_df['QrSp2'] = vol_df['Qr']*misc_df['S2mean']

        if old_dt:
            smean2_net = misc_df['Smean'].to_numpy()**2 * misc_df['V'].to_numpy()
            dsmean2_net_dt = (smean2_net[2:] - smean2_net[:-2]) / (2*86400)
        else:
            dsmean2_net_dt = dsmean2v_dt_lp

        if old_dt:
            sp2_df.loc[1:-1,'dSp2net_dt'] = ds2net_dt - dsmean2_net_dt
        else:
            sp2_df.loc[:,'dSp2net_dt'] = ds2net_dt - dsmean2_net_dt

        sp2_df['Mix_resolved'] = salt2_df['Mix_resolved']
        sp2_df['Mix_numerical'] = sp2_df['dSp2net_dt'] - sp2_df['QSp2in'] - sp2_df['QSp2out'] - sp2_df['QrSp2'] - sp2_df['Mix_resolved']
        sp2_df['Mixing'] = sp2_df['Mix_resolved'] + sp2_df['Mix_numerical']

        # Normalized S'^2 budget
        QSS = vol_df['Qr'].mean() * misc_df['Sin'] * misc_df['Sin']
        sp2n_df = sp2_df.copy()
        for cn in sp2n_df.columns:
            sp2n_df[cn] = sp2n_df[cn]/QSS

        # Mixedness budget
        m_df = pd.DataFrame(index=indall)
        m_df['dmnet_dt'] = Socn*salt_df['dSnet_dt'] - salt2_df['dS2net_dt']
        m_df['Qmin'] = vol_df['Qin']*(Socn*misc_df['Sin'] - misc_df['S2in'])
        m_df['Qmout'] = vol_df['Qout']*(Socn*misc_df['Sout'] - misc_df['S2out'])
        m_df['Mixing'] = -salt2_df['Mixing']
        m_df['Error'] = m_df['dmnet_dt'] - m_df['Qmin'] - m_df['Qmout'] - m_df['Mixing']

        # normalized mixedness
        mn_df = m_df * (365 * 86400) / (misc_df['V'].mean() * (Socn/2) * (Socn/2))

        # make sure everything is numeric
        for cn in misc_df.columns:
            misc_df[cn] = pd.to_numeric(misc_df[cn])
        for cn in vol_df.columns:
            vol_df[cn] = pd.to_numeric(vol_df[cn])
        for cn in salt_df.columns:
            salt_df[cn] = pd.to_numeric(salt_df[cn])
        for cn in salt2_df.columns:
            salt2_df[cn] = pd.to_numeric(salt2_df[cn])
        for cn in salt_qe_df.columns:
            salt_qe_df[cn] = pd.to_numeric(salt_qe_df[cn])
        for cn in sp2_df.columns:
            sp2_df[cn] = pd.to_numeric(sp2_df[cn])
        for cn in sp2n_df.columns:
            sp2n_df[cn] = pd.to_numeric(sp2n_df[cn])
        for cn in m_df.columns:
            m_df[cn] = pd.to_numeric(m_df[cn])
        for cn in m_df.columns:
            mn_df[cn] = pd.to_numeric(mn_df[cn])

        # save DataFrames to pickle files
        misc_df.to_pickle(outdir + 'misc_df_' + year_str + '_' + which_vol.replace(' ','_') + '.p')
        vol_df.to_pickle(outdir + 'vol_df_' + year_str + '_' + which_vol.replace(' ','_') + '.p')
        salt_df.to_pickle(outdir + 'salt_df_' + year_str + '_' + which_vol.replace(' ','_') + '.p')
        salt2_df.to_pickle(outdir + 'salt2_df_' + year_str + '_' + which_vol.replace(' ','_') + '.p')
        sp2_df.to_pickle(outdir + 'sp2_df_' + year_str + '_' + which_vol.replace(' ','_') + '.p')
        sp2n_df.to_pickle(outdir + 'sp2n_df_' + year_str + '_' + which_vol.replace(' ','_') + '.p')
        m_df.to_pickle(outdir + 'm_df_' + year_str + '_' + which_vol.replace(' ','_') + '.p')

        # accumulate error statistics (all %)
        err_df_vol.loc[year, which_vol] = vol_rel_err*100
        err_df_salt.loc[year, which_vol] = salt_rel_err*100
        err_df_salt2.loc[year, which_vol] = salt2_rel_err*100

        # debugging: get net salt and salt-squared flux as a check
        if False:
            salt_df['QSnet'] = 0
            salt2_df['QS2net'] = 0
            tt0 = time()
            for sect_name in sect_sign_dict.keys():
                print('Debugging - processing ' + sect_name)
                sys.stdout.flush()
                sect_sign = sect_sign_dict[sect_name]
                ext_fn = indir0 + 'extractions/' + sect_name + '.nc'
                ext = nc.Dataset(ext_fn)
                ext_q = sect_sign * ext['q'][:]
                ext_salt = ext['salt'][:]
                xqs = ((ext_q * ext_salt).sum(axis=2)).sum(axis=1)
                xqs2 = ((ext_q * ext_salt * ext_salt).sum(axis=2)).sum(axis=1)
                xqs_lp = zfun.filt_godin(xqs)
                xqs2_lp = zfun.filt_godin(xqs2)
                pad = 36
                xqs_lp = xqs_lp[pad:-(pad+1):24]
                xqs2_lp = xqs2_lp[pad:-(pad+1):24]
                salt_df['QSnet'] += xqs_lp
                salt2_df['QS2net'] += xqs2_lp
                print('  -- took %0.2f sec' % (time()-tt0))
                sys.stdout.flush()
            salt_df['Error_alt'] = salt_df['dSnet_dt'] - salt_df['QSnet']
            salt2_df['Mix_numerical_alt'] = salt2_df['dS2net_dt'] - salt2_df['QS2net'] - salt2_df['Mix_resolved']

        # --------------- Volume and Salt Budgets --------------------------------------
        fig1 = plt.figure(figsize=(16,7))

        ax = fig1.add_subplot(121)
        vol_df[['dV_dt','Qin','Qout', 'Qr','Error']].plot(ax=ax, grid=True).legend(loc='upper right')
        ax.set_title('Volume Budget (m3/s)')
        ax.text(.05,.9, 'Mean Error / Mean Qr = %0.2f%%' % (vol_rel_err*100), transform=ax.transAxes, fontsize=14)

        ax = fig1.add_subplot(122)
        #salt_df[['dSnet_dt','QSin','QSout','Error']].plot(ax=ax, grid=True).legend(loc='upper right')
        salt_df.plot(ax=ax, grid=True).legend(loc='upper right')
        ax.set_title(year_str + ' ' + which_vol + ' Salt Budget (g/kg m3/s)')
        ax.text(.05,.9, 'Mean Error / Mean QSin = %0.2f%%' % (salt_rel_err*100), transform=ax.transAxes, fontsize=14)

        if save_figs:
            outname1 = outdir + 'vol_salt_budget_' + year_str + '_' + which_vol.replace(' ','_')
            fig1.savefig(outname1 + '.png')

        # --------------- Salinity-squared Budget ------------------------------------------
        fig2 = plt.figure(figsize=(16,7))

        ax = fig2.add_subplot(121)
        #salt2_df[['dS2net_dt','QS2in','QS2out','Mix_resolved','Mix_numerical']].plot(ax=ax, grid=True).legend(loc='upper right')
        salt2_df.plot(ax=ax, grid=True).legend(loc='upper right')
        ax.set_title(year_str + ' ' + which_vol + ' Salt2 Budget (g2/kg2 m3/s)')

        ax = fig2.add_subplot(122)
        salt2_df[['Mix_resolved','VMix','HMix','Mix_numerical']].plot(ax=ax, grid=True).legend(loc='upper right')
        ax.set_title(year_str + ' ' + which_vol + ' Salt2 Budget (g2/kg2 m3/s)')
        ax.text(.05,.9, 'Mean Error / Mean Mix_resolved = %0.2f%%' % (salt2_rel_err*100), transform=ax.transAxes, fontsize=14)

        if save_figs:
            outname2 = outdir + 'salt2_budget_' + year_str + '_' + which_vol.replace(' ','_')
            fig2.savefig(outname2 + '.png')

        # --------------- S'^2 squared Budget ----------------------------------------
        fig3 = plt.figure(figsize=(12,7))

        ax = fig3.add_subplot(111)
        sp2n_df[['dSp2net_dt', 'QSp2in', 'QSp2out', 'QrSp2', 'Mixing', 'Mix_resolved']].plot(ax=ax, grid=True).legend(loc='upper right')
        ax.set_title(year_str + ' ' + which_vol + ' Sp2 Budget Normalized')
        ax.set_ylim(-4,4)

        if save_figs:
            outname3 = outdir + 'sp2n_budget_' + year_str + '_' + which_vol.replace(' ','_')
            fig3.savefig(outname3 + '.png')

        # --------------- Mixedness Budget ------------------------------------
        fig4 = plt.figure(figsize=(12,7))

        ax = fig4.add_subplot(211)
        m_df[['dmnet_dt', 'Qmin', 'Qmout', 'Mixing', 'Error']].plot(ax=ax, grid=True).legend(loc='upper right')
        ax.set_title(year_str + ' ' + which_vol + ' Mixedness Budget')

        ax = fig4.add_subplot(212)
        misc_df['Qprism'].plot(ax=ax, grid=True)
        #ax.set_ylim(-4,4)

        if save_figs:
            outname4 = outdir + 'm_budget_' + year_str + '_' + which_vol.replace(' ','_')
            fig4.savefig(outname4 + '.png')

        # --------------- Normalized Mixedness Budget ------------------------------------
        fig5 = plt.figure(figsize=(12,7))

        ax = fig5.add_subplot(211)
        mn_df[['dmnet_dt', 'Qmin', 'Qmout', 'Mixing', 'Error']].plot(ax=ax, grid=True).legend(loc='upper right')
        ax.set_title(year_str + ' ' + which_vol + ' Normalized Mixedness Budget')
        ax.set_ylabel('V*(Socn/2)^2 per year')

        ax = fig5.add_subplot(212)
        misc_df['Qprism'].plot(ax=ax, grid=True)
        #ax.set_ylim(-4,4)

        if save_figs:
            outname5 = outdir + 'm_budget_' + year_str + '_' + which_vol.replace(' ','_')
            fig5.savefig(outname5 + '.png')

        # --------------- Checking ---------------------------------------------------
        fig6 = plt.figure(figsize=(12,6))

        ax = fig6.add_subplot(111)
        salt2_df[['Mixing']].plot(ax=ax, grid=True, label='from salt2',style='-r').legend(loc='upper right')
        sp2_df[['Mixing']].plot(ax=ax, grid=True, label='from sp2',style='-b').legend(loc='upper right')
        ax.set_title(year_str + ' ' + which_vol + ' Mixing (g2/kg2 m3/s)')
        ax.text(.05, .9, 'Mean Mixing from salt2 = %.5e' % (salt2_df['Mixing'].mean()), transform=ax.transAxes, c='r')
        ax.text(.05, .8, 'Mean Mixing from sp2 = %.5e' % (sp2_df['Mixing'].mean()), transform=ax.transAxes, c='b')
        ax.text(.05, .2, 'Ratio = %.3f' % (salt2_df['Mixing'].mean()/sp2_df['Mixing'].mean()), transform=ax.transAxes, c='purple')

plt.show()

max_stars_repo_name: parkermac/LiveOcean

max_issues_repo_name: parkermac/LiveOcean

max_forks_repo_name: parkermac/LiveOcean


text:
#-----------------------------------------------------------------------------
# Copyright (c) 2014, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

from yt.visualization.volume_rendering.theia.source import TheiaSource
from yt.visualization.volume_rendering.theia.camera import Camera

class TheiaScene:
    r"""
        TheiaScene

        This is the highest level entry point into Theia.
        This class is a container for a Theia source and
        a camera for controlling a view onto the source.

        Call :
  	    update()
        To perform the volume rendering.
        Call :
            get_results()
        To get the resulting image array.

        Parameters
        ----------
        ts : TheiaSource
            The container for the raycasting algorithm and the volumetric data
        volume :
            If a TheiaSource is not specified a volme can be used to initialize
            the scene. However, rendering will only occur when a raycaster
            is also provided.

        camera :
            A camera for controling view point. Must be able to provide a 4x4
            matrix to the raycaster via get_matrix()
            If no camera is provided the default Theia Camera is used.

        Example:

        from yt.visualization.volume_rendering.theia.scene import TheiaScene
        from yt.visualization.volume_rendering.algorithms.front_to_back import FrontToBackRaycaster
        import numpy as np

        #load 3D numpy array of float32
        volume = np.load("/home/bogert/log_densities_1024.npy")

        scene = TheiaScene( volume = volume, raycaster = FrontToBackRaycaster() )

        scene.camera.rotateX(1.0)
        scene.update()

        surface = scene.get_results()
        #now do something with surface

    """
    def __init__(self, ts = None, camera = None, volume = None, raycaster = None) :
        #Values to control screen resolution
        if ts != None :
            self.source = ts
        else :
            self.source = TheiaSource(volume  = volume, raycaster = raycaster)

        if camera == None :
              self.camera = Camera()
        else :
              self.camera = camera

    def get_raycaster(self):
        """
            This is for programmers who would prefer calling :
                scene.get_raycaster()
            instead of :
                scene.source.raycaster
	"""
        return self.source.raycaster

    def update(self):
        """
            This will send the camera's modelview matrix to
            the raycaster and call the raycaster's cast function
            to perform the volume rendering.
        """
        self.source.update(self.camera.get_matrix())

    def get_results(self):
        """
            This will cause the results of the volume rendering to be
            copied back from the GPU and then returned as a 2d numpy array.
        """
        return self.source.get_results()


text:
import numpy as np
import torch
from torch import nn

from .base import Flow

class Planar(Flow):
    """
    Planar flow as introduced in arXiv: 1505.05770
        f(z) = z + u * h(w * z + b)
    """

    def __init__(self, shape, act="tanh", u=None, w=None, b=None):
        """
        Constructor of the planar flow
        :param shape: shape of the latent variable z
        :param h: nonlinear function h of the planar flow (see definition of f above)
        :param u,w,b: optional initialization for parameters
        """
        super().__init__()
        lim_w = np.sqrt(2. / np.prod(shape))
        lim_u = np.sqrt(2)

        if u is not None:
            self.u = nn.Parameter(u)
        else:
            self.u = nn.Parameter(torch.empty(shape)[None])
            nn.init.uniform_(self.u, -lim_u, lim_u)
        if w is not None:
            self.w = nn.Parameter(w)
        else:
            self.w = nn.Parameter(torch.empty(shape)[None])
            nn.init.uniform_(self.w, -lim_w, lim_w)
        if b is not None:
            self.b = nn.Parameter(b)
        else:
            self.b = nn.Parameter(torch.zeros(1))

        self.act = act
        if act == "tanh":
            self.h = torch.tanh
        elif act == "leaky_relu":
            self.h = torch.nn.LeakyReLU(negative_slope=0.2)
        else:
            raise NotImplementedError('Nonlinearity is not implemented.')

    def forward(self, z):
        lin = torch.sum(self.w * z, list(range(1, self.w.dim()))) + self.b
        if self.act == "tanh":
            inner = torch.sum(self.w * self.u)
            u = self.u + (torch.log(1 + torch.exp(inner)) - 1 - inner) * self.w / torch.sum(self.w ** 2)
            h_ = lambda x: 1 / torch.cosh(x) ** 2
        elif self.act == "leaky_relu":
            inner = torch.sum(self.w * self.u)
            u = self.u + (torch.log(1 + torch.exp(inner)) - 1 - inner) * self.w / torch.sum(
                self.w ** 2)  # constraint w.T * u neq -1, use >
            h_ = lambda x: (x < 0) * (self.h.negative_slope - 1.0) + 1.0

        z_ = z + u * self.h(lin.unsqueeze(1))
        log_det = torch.log(torch.abs(1 + torch.sum(self.w * u) * h_(lin)))
        return z_, log_det

    def inverse(self, z):
        if self.act != "leaky_relu":
            raise NotImplementedError('This flow has no algebraic inverse.')
        lin = torch.sum(self.w * z, list(range(2, self.w.dim())), keepdim=True) + self.b
        inner = torch.sum(self.w * self.u)
        a = ((lin + self.b) / (1 + inner) < 0) * (self.h.negative_slope - 1.0) + 1.0  # absorb leakyReLU slope into u
        u = a * (self.u + (torch.log(1 + torch.exp(inner)) - 1 - inner) * self.w / torch.sum(self.w ** 2))
        z_ = z - 1 / (1 + inner) * (lin + u * self.b)
        log_det = -torch.log(torch.abs(1 + torch.sum(self.w * u)))
        if log_det.dim() == 0:
            log_det = log_det.unsqueeze(0)
        if log_det.dim() == 1:
            log_det = log_det.unsqueeze(1)
        return z_, log_det


text:
import numpy as np

class Layer:
    def __init__(self, connections, neurons, activationFunction):
        self.activationFunction = activationFunction
        self.w = np.random.rand(connections, neurons) * 2 - 1
        self.b = np.random.rand(1, neurons) * 2 - 1