text:
from __future__ import annotations
from typing import Any, Callable, Optional, cast

from kit import gcopy, implements
import numpy as np
import pytorch_lightning as pl
from pytorch_lightning.callbacks.progress import ProgressBar
from pytorch_lightning.utilities.types import STEP_OUTPUT
from torch import Tensor, nn, optim
from torch.utils.data import DataLoader

from extinct.components.datamodules import DataBatch, Stage, VisionDataModule
from extinct.components.models.base import ModelBase
from extinct.utils import cosine_scheduler, get_params_groups

from . import vit
from .eval import KNN, DatasetEncoderRunner, DINOLinearClassifier, EvalMethod
from .loss import DINOLoss

__all__ = ["DINO"]

from extinct.components.callbacks.dino_updates import MeanTeacherWeightUpdate

from .models import MultiCropNet

class DINO(ModelBase):
    _loss_fn: DINOLoss
    student: MultiCropNet
    teacher: MultiCropNet
    eval_trainer: pl.Trainer
    lr_schedule: np.ndarray
    wd_schedule: np.ndarray
    momentum_schedule: np.ndarray
    datamodule: VisionDataModule
    eval_clf: DINOLinearClassifier | KNN

    def __init__(
        self,
        lr: float = 5.0e-4,
        warmup_iters: int = 10,
        weight_decay: float = 4.0e-2,
        min_lr: float = 1.0e-6,
        weight_decay_end: float = 0.4,
        freeze_last_layer: int = 1,
        arch: vit.VitArch = vit.VitArch.small,
        patch_size: int = 16,
        out_dim: int = 65_536,
        norm_last_layer: bool = True,
        use_bn_in_head: bool = False,
        momentum_teacher: float = 0.996,
        teacher_temp: float = 0.04,
        warmup_teacher_temp_iters: int = 30,
        eval_method: EvalMethod = EvalMethod.lin_clf,
        num_eval_blocks: int = 1,
        lr_eval: float = 1.0e-4,
        lin_clf_epochs: int = 100,
        batch_size_eval: Optional[int] = None,
        max_steps: int = -1,
        dm_batch_size: int = -1,
        local_crops_number: int = -1,
    ) -> None:
        """
        Args:
            num_eval_blocks: Concatenate [CLS] tokens for the `n` last blocks.
            We use `n=4` when evaluating ViT-Small and `n=1` with ViT-Base.
        """
        super().__init__()
        self.learning_rate = lr
        self.num_eval_blocks = num_eval_blocks
        self.warmup_iters = warmup_iters
        self.weight_decay = weight_decay
        self.weight_decay_end = weight_decay_end
        self.min_lr = min_lr
        self.freeze_last_layer = freeze_last_layer
        self.patch_size = patch_size
        self.out_dim = out_dim
        self.norm_last_layer = norm_last_layer
        self.use_bn_in_head = use_bn_in_head
        self.momentum_teacher = momentum_teacher
        self.teacher_temp = teacher_temp
        self.warmup_teacher_temp_iters = warmup_teacher_temp_iters
        self.lr_eval = lr_eval
        self.eval_method = eval_method
        self.lin_clf_epochs = lin_clf_epochs
        self.batch_size_eval = batch_size_eval

        self.weight_callback = MeanTeacherWeightUpdate(
            max_steps=max_steps, initial_tau=self.momentum_teacher
        )

        self._arch_fn = cast(
            Callable[[int], vit.VisionTransformer], getattr(vit, f"vit_{arch.name}")
        )

        self._loss_fn = DINOLoss(
            out_dim=self.out_dim,
            warmup_teacher_temp=self.teacher_temp,
            teacher_temp=self.teacher_temp,
            warmup_teacher_temp_iters=min(max_steps - 1, self.warmup_teacher_temp_iters),
            total_iters=max_steps,  # type: ignore
            num_crops=local_crops_number + 2,
        )

        self.lr_schedule = cosine_scheduler(
            base_value=self.learning_rate * dm_batch_size / 256.0,  # linear scaling rule
            final_value=self.min_lr,
            total_iters=max_steps,  # type: ignore
            warmup_iters=min(max_steps - 1, self.warmup_iters),
        )
        self.wd_schedule = cosine_scheduler(
            base_value=self.weight_decay,
            final_value=self.weight_decay_end,
            total_iters=max_steps,
        )

        self.student = MultiCropNet(
            arch_fn=self._arch_fn,
            patch_size=self.patch_size,
            norm_last_layer=self.norm_last_layer,
            use_bn_in_head=self.use_bn_in_head,
            out_dim=self.out_dim,
        )
        self.teacher = MultiCropNet(
            arch_fn=self._arch_fn,
            patch_size=self.patch_size,
            norm_last_layer=True,
            use_bn_in_head=self.use_bn_in_head,
            out_dim=self.out_dim,
        )

    @property
    def enc(self) -> vit.VisionTransformer:
        # We define an encoder-extracting method for consistency with the other models -
        # fit_and_test, for instance, expects a model to comprise of two parts: enc and clf.
        return self.student.backbone

    def _encode_dataset(self, stage: Stage) -> DataBatch:
        # It's not strictly necessary to disable shuffling but pytorch-lightning complains if its
        # enabled during 'testing'
        dl_kwargs = dict(eval_=True) if stage == "train" else {}
        # Sampler needs to be set to None, meaning the default sequential/batch sampler combination
        # is used, so that the full dataset is encoded (with no duplicates)
        dataloader = cast(DataLoader, getattr(self.datamodule, f"{stage}_dataloader")(**dl_kwargs))
        # Encode the dataset
        dataset_encoder = DatasetEncoderRunner(model=self.student.backbone)
        self.eval_trainer.test(
            dataset_encoder,
            test_dataloaders=dataloader,
            verbose=False,
        )
        # Extract the encodings/associated labels from the dataset encoder
        return dataset_encoder.encoded_dataset

    @implements(ModelBase)
    def build(self, datamodule: VisionDataModule, trainer: pl.Trainer) -> None:
        self.datamodule = datamodule
        self.eval_trainer = gcopy(
            trainer,
            deep=True,
            max_epochs=self.lin_clf_epochs,
            max_steps=None,
        )
        bar = ProgressBar()
        bar._trainer = self.eval_trainer
        self.eval_trainer.callbacks = [bar]

    @implements(pl.LightningModule)
    def configure_optimizers(self) -> optim.Optimizer:
        return optim.AdamW(
            get_params_groups(self.student), lr=self.learning_rate, weight_decay=self.weight_decay
        )

    def _cancel_gradients_last_layer(self, train_itr: int) -> None:
        if train_itr >= self.freeze_last_layer:
            return
        for n, p in self.student.named_parameters():
            if "last_layer" in n:
                p.grad = None

    def _get_loss(self, batch: DataBatch, batch_idx: int) -> Tensor:
        teacher_output = self.teacher(
            batch.x[:2]
        )  # only the 2 global views pass through the teacher
        student_output = self.student(batch.x)
        return self._loss_fn(student_output, teacher_output, batch_idx)

    @implements(pl.LightningModule)
    def training_step(self, batch: DataBatch, batch_idx: int) -> Tensor:
        for i, param_group in enumerate(self.trainer.optimizers[0].param_groups):
            param_group["lr"] = self.lr_schedule[batch_idx]
            if i == 0:  # only the first group is regularized
                param_group["weight_decay"] = self.wd_schedule[batch_idx]

        return self._get_loss(batch=batch, batch_idx=batch_idx)

    def on_train_batch_end(
        self, outputs: STEP_OUTPUT, batch: Any, batch_idx: int, dataloader_idx: int
    ) -> None:
        self.weight_callback.on_train_batch_end(
            self.trainer, self, outputs, batch, batch_idx, dataloader_idx
        )

    @implements(pl.LightningModule)
    def optimizer_step(
        self,
        epoch: int,
        batch_idx: int,
        optimizer: optim.Optimizer,
        optimizer_idx: int,
        optimizer_closure: Callable | None,
        on_tpu: bool,
        using_native_amp: bool,
        using_lbfgs: bool,
    ) -> None:
        # Keep the output layer fixed until the epoch-threshold has been reached
        # Typically doing so during the first epoch helps training.
        self._cancel_gradients_last_layer(train_itr=batch_idx)
        # Update the student's parameters using the DINO loss
        super().optimizer_step(
            epoch=epoch,
            batch_idx=batch_idx,
            optimizer=optimizer,
            optimizer_idx=optimizer_idx,
            optimizer_closure=optimizer_closure,
            on_tpu=on_tpu,
            using_native_amp=using_native_amp,
            using_lbfgs=using_lbfgs,
        )

    @implements(ModelBase)
    def _inference_step(self, batch: DataBatch, stage: Stage) -> dict[str, Any]:
        return self.eval_clf._inference_step(batch=batch, stage=stage)

    @implements(ModelBase)
    def _inference_epoch_end(
        self, output_results: list[dict[str, Tensor]], stage: Stage
    ) -> dict[str, Any]:
        return self.eval_clf._inference_epoch_end(output_results=output_results, stage=stage)

    @implements(nn.Module)
    def forward(self, x: Tensor) -> Tensor:
        return self.student(x)

max_stars_repo_name: olliethomas/extinct-dino

max_issues_repo_name: olliethomas/extinct-dino

max_forks_repo_name: olliethomas/extinct-dino


text:
# ----------------------------------------------------------------------------
# This software is in the public domain, furnished "as is", without technical
# support, and with no warranty, express or implied, as to its usefulness for
# any purpose.
#
# BOIVerifyBiasCorr2 - Version 2.0.5
#
#    Run the BOIVerifyBiasCorr tool
#
# Author: Tim Barker - SOO BOI
#   2005/07/01 - version 0.0 - original implementation
#   2005/07/29 - version 0.1 - update to grid database structure
#   2006/11/06 - version 1.0 - Make only one procedure!  And have parms
#                    specified via the file name FILE (constant below).
#   2007/10/25 - version 2.0 - Everything now in procedure - for simpler
#                    maintenance and less tool/procedure complexity.
#                    Tones down regression estimate when forecast value is
#                    well outside recent forecasts - falling back to simple
#                    average error over past VERPERIOD days. Obs database
#                    must now be specified. Better logging. Better memory
#                    usage.
#   2008/05/28 - version 2.0.5 - bugfix for basetime/starttime when an
#                    archived grid is constant.
#
# ----------------------------------------------------------------------------
#
#     SOFTWARE HISTORY
#
#    Date            Ticket#       Engineer       Description
#    ------------    ----------    -----------    --------------------------
#    11/21/13        16770         ryu            Change name of temporary files
#                                                 for dual domain.
#=============================================================================

##
# This is an absolute override file, indicating that a higher priority version
# of the file will completely replace a lower priority version of the file.
##

#
#  Do not show this in any menu.  Should only be run via runProcedure after
#  putting the force flag and model to correct in /tmp/<siteId>_FILE
#
#MenuItems = ["Verify"]
#
MenuItems = ["None"]
#
#  Constants
#
FILE="BOIVerifyBiasCorr.txt"
HOURSECS=60*60
DAYSECS=24*HOURSECS
#
#  imports
#
from numpy import *
import os, time
import SmartScript,BOIVerifyUtility
import AbsTime
import Exceptions
#
#
DSPATTR_PORT = 0  # TODO - display attributes not working the same way
#
#
class Procedure (SmartScript.SmartScript):
   def __init__(self, dbss):
      SmartScript.SmartScript.__init__(self, dbss)
      self._dbss=dbss

   def execute(self, editArea):
      (gyea,gmon,gday,ghou,gmin,gsec,gwda,gyda,gdst)=time.gmtime()
      self.VU=BOIVerifyUtility.BOIVerifyUtility(self._dbss, None)
      self.VU.logMsg("BOIVerifyBiasCorr Procedure Start")

      #
      #
      #  BIASCORR_DAYS configuration should hold number of days usually used
      #  in the bias correction.  If it doesn't...we default to 30
      #
      VERPERIOD=self.VU.getCFG('BIASCORR_DAYS')
      if VERPERIOD is None:
         VERPERIOD=30
      #
      #  BIASCORR_MINDAYS configuration should hold minimum number of days
      #  needed to do a regression.  If it doesn't...we default to 14
      #
      MINDAYS=self.VU.getCFG('BIASCORR_MINDAYS')
      if MINDAYS is None:
         MINDAYS=14
      #
      #  BIASCORR_EXTRAPOLATE_PERCENTAGE configuration should hold the
      #  percentage of range of recent forecasts that we will extrapolate
      #  with regression.  Outside of this...drop back toward using
      #  only average bias
      #
      FUZZ=self.VU.getCFG('BIASCORR_EXTRAPOLATE_PERCENTAGE')
      if FUZZ is None:
         FUZZ=10
      fuzz=FUZZ/100
      #
      #  Thes are the Parms we currently know how to bias correct
      #
      parmList=["T","MaxT","MinT","RH","MaxRH","MinRH"]
      self.transformParms=["RH","MaxRH","MinRH"]
      #
      #  Read the force flag, model, and obsmodel from the FILE file
      #
      force="0"
      model=""
      obsmodel=""
      filename="/tmp/%s_%s"% (self.getSiteID(), FILE)
      if ((os.path.exists(filename))and(os.path.isfile(filename))):
         try:
            infile=file(filename,"r")
            inlines=infile.readlines()
            infile.close()
            os.remove(filename)
            if len(inlines)>0:
               line=inlines[0]
               force=line.strip()
            if len(inlines)>1:
               line=inlines[1]
               model=line.strip()
            if len(inlines)>2:
               line=inlines[2]
               obsmodel=line.strip()
         except:
            self.VU.logMsg("Difficulty reading %s"%filename)
      ##
      ##  testing
      ##
      ##model="NAM12"
      ##obsmodel="Obs"
      ##force="1"
      ##self.VU.setDebug(10)
      #
      #  If model is blank - give error and stop
      #
      if (model==""):
          self.VU.logMsg("No model to correct - cannot continue")
          return
      if (obsmodel==""):
          self.VU.logMsg("No obs model specified - cannot continue")
          return
      #
      #  if obsmodel is not a valid obsmodel - stop
      #
      if obsmodel not in self.VU.getCFG('OBSMODELS'):
         self.VU.logMsg("%s is not a valid name for observations"%obsmodel)
         return
      #
      #  Get corresponding (non-bc) model run with same basetime
      #
      mutableModel=self.mutableID().modelName()
      modeltime=self.mutableID().modelTime().unixTime()
      mtime=time.strftime("%Y%m%d_%H%M",time.gmtime(modeltime))
      foundit=0
      for run in xrange(0,-10,-1):
         dbid=self.findDatabase(model,version=run)
         if dbid is not None:
            inputtime=dbid.modelTime().unixTime()
            if inputtime==modeltime:
               foundit=1
               break
      if foundit==0:
         self.VU.logMsg("%s run from %s could not be found"%(model,mtime))
         self.VU.logMsg("Aborting!")
         return
      (myr,mmo,mdy,mhr,mmi,mse,mwd,myd,mds)=time.gmtime(modeltime)
      yesterday=modeltime-DAYSECS
      modelCycle=mhr
      bigtr=self.createTimeRange(0,500,"Database",dbid)
      self.VU.logMsg("for %s at %2.2dZ"%(model,modelCycle))
      #
      #  Loop over parms
      #
      for parm in parmList:
         self.VU.logMsg("Working on %s"%parm)
         (self.parmUnits,self.parmPrecision,self.parmMinval,self.parmMaxval,
          self.parmColorTable,self.parmDisplayMinval,
          self.parmDisplayMaxval)=self.getParmInfo(model,parm)
         obsparm=self.VU.getObsParm(parm)
         #
         #  get corresponding pairs of observed/fcst grids for last VERPERIOD days
         #
         dateStyle="Verifying on"
         dateType="Period Length"
         self.VU.logMsg("    getting cases",5)
         cases=self.VU.getCases(parm,model,obsparm,obsmodel,dateStyle,dateType,
                                fromDay=yesterday,numDays=VERPERIOD,
                                cycles=mhr,fhrEnd=240)
         caseInfo=cases[model]
         del cases
         fhrInfo={}
         for key in caseInfo.keys():
            (basestr,startstr,endstr)=key.split(",")
            basetime=int(basestr)
            starttime=int(startstr)
            fhrInfo[key]=(starttime-basetime)/HOURSECS
         self.VU.logMsg("    done with getting cases",5)
         #
         #  Loop over all forecast grids
         #
         try:
            gridInfoList=self.getGridInfo(dbid,parm,"SFC",bigtr)
         except Exception, e:
            self.VU.logMsg(str(e))
            continue

         for gridInfo in gridInfoList:
            starttime=gridInfo.gridTime().startTime().unixTime()
            endtime=gridInfo.gridTime().endTime().unixTime()
            forecastHour=int((starttime-modeltime)/3600.0)
            endHour=int((endtime-modeltime)/3600.0)
            self.VU.logMsg("  %d-hr forecast"%forecastHour)
            #
            #  Check to see if guess grid is newer than output grid
            #  (with 10 minute overlap...)
            #
            origtr=self.createTimeRange(forecastHour,endHour,"Database",dbid)
            guesstime=self.getLastUpdate(dbid,parm,"SFC",origtr)
            outtime=self.getLastUpdate(mutableModel,parm,"SFC",origtr)
            if ((guesstime<=(outtime-(10*60)))and(force=="0")):
               self.VU.logMsg("    no new info - skipping")
               continue
            #
            #  Read the guess grid
            #
            inputGrid=self.getGrids(dbid,parm,"SFC",origtr,noDataError=0,cache=0)
            if inputGrid is None:
               self.VU.logMsg("Could not read input grid")
               continue
            if parm in self.transformParms:
               inputGrid=self.Ztrans(inputGrid)
            #
            #  Get old forecasts for this forecast hour
            #
            t1=time.time()
            result=self.getOldForecasts(caseInfo,fhrInfo,parm,model,obsparm,
                                        obsmodel,forecastHour,self.parmMinval,
                                        self.parmMaxval)
            t2=time.time()
            self.VU.logMsg("    getting old forecasts took %6.2f"%(t2-t1),5)
            (obsList,fcstList)=result
            numGrids=len(obsList)
            del result
            #
            #  Do not calculate any grids if no previous observations/forecasts
            #
            if numGrids<1:
               continue
            #
            #
            if numGrids<MINDAYS:
               self.VU.logMsg("    Only %d old forecast/observation grids"%numGrids)
               self.VU.logMsg("    Not enough to make regression - skipping")
               continue
            #
            #  Turn list of obs grids and fcst grids into large array
            #
            obsGrids=array(obsList)
            fcstGrids=array(fcstList)
            del obsList
            del fcstList
            #
            #  Get average errors
            #
            avgFcst=add.reduce(fcstGrids)/numGrids
            avgErr=add.reduce(fcstGrids-obsGrids)/numGrids
            maxFcst=maximum.reduce(fcstGrids)
            minFcst=minimum.reduce(fcstGrids)
            #
            #  If more than MINDAYS old forecast/observation pairs...add
            #  regression information
            #
            if numGrids>=MINDAYS:
               self.VU.logMsg("    %d old obs/fcst grids used for regression"%numGrids,2)
               #
               #
               #  Linear regression of forecast anomalies to non-average
               #  forecast errors.
               #
               t1=time.time()
               result=self.getRegression(avgFcst,avgErr,obsGrids,fcstGrids)
               t2=time.time()
               self.VU.logMsg("    getting Regression took %6.3f"%(t2-t1),5)
               (correlation,stdind,stddep)=result
               del result
               del obsGrids
               del fcstGrids
               #
               #  Create the regressed part of error grid from the linear regression
               #
               regErr=(((correlation*stddep)/stdind)*(inputGrid-avgFcst))
               del avgFcst
               del stddep
               del stdind
               del correlation
            else:
               self.VU.logMsg("   %d old obs/fcst grids - used plain bias"%numGrids,2)
               regErr=self.empty()
            #
            #  Make a multiplier for the regressed error.  Normally 1.0...but
            #  when the forecast is more than fuzz% beyond the range of forecasts
            #  in the training period...start cutting back correction amount
            #  until getting back to zero regressed amount at (2*fuzz)% beyond.
            #
            multiplier=self.newGrid(1)
            fuzzGrid=maximum((maxFcst-minFcst)*fuzz,0.01) # dont have zero
            max1=maxFcst+fuzzGrid
            multiplier=where(greater(inputGrid,max1),1.0-((inputGrid-max1)/fuzzGrid),multiplier)
            min1=minFcst-fuzzGrid
            multiplier=where(less(inputGrid,min1),1.0-((min1-inputGrid)/fuzzGrid),multiplier)
            multiplier[less(multiplier,0.0)] = 0.0
            del min1
            del max1
            del minFcst
            del maxFcst
            del fuzzGrid
            if self.VU.getDebug>=1:
               count=count_nonzero(less(multiplier,0.98))
               count2=count_nonzero(less(multiplier,0.02))
               if count>0:
                  self.VU.logMsg("    %d fcst points are outliers and regression was reduced"%count,1)
                  if count2>0:
                     self.VU.logMsg("       %d points so far out that only average error used"%count2,1)
            #
            #  Correct the forecast with the predicted error
            #
            corrected=inputGrid-(avgErr+(regErr*multiplier))
            del regErr
            del avgErr
            del inputGrid
            del multiplier
            #
            #  For Transformed Parms (like RH) transform it back to the real variable
            #
            if parm in self.transformParms:
               corrected=self.Rtrans(corrected)
            self.VU.logMsg("    done with making corrected grid",5)
            #
            #  clip to legal values and save
            #
            corrected=clip(corrected,self.parmMinval,self.parmMaxval)
            self.createGrid(mutableModel,parm,"SCALAR",corrected,origtr,
                                 parm,None,self.parmPrecision,self.parmMinval,
                                 self.parmMaxval,self.parmUnits)
            del corrected
            self.VU.logMsg("    mem:%d res:%d"%(memory(),resident()),10)
         self.saveGrid(mutableModel,parm)
      #
      #  Fix issues that can arise via doing each parm independently
      #
      #
      #  Make sure MaxT is as high as the max of all T grids
      #
      self.VU.logMsg("Checking for hourly T grids higher than MaxT")
      gridInfoList=self.getGridInfo(mutableModel,"MaxT","SFC",bigtr)
      for gridInfo in gridInfoList:
         newtr=gridInfo.gridTime()
         maxtgrid=self.getGrids(mutableModel,"MaxT","SFC",newtr,noDataError=0,cache=0)
         if maxtgrid is not None:
            maxoftgrid=self.getGrids(mutableModel,"T","SFC",newtr,mode="Max",noDataError=0,cache=0)
            if maxoftgrid is not None:
               maxt=maximum(maxoftgrid,maxtgrid)
               changed=count_nonzero(greater(maxt,maxtgrid))
               if changed>0:
                  fhr=int((newtr.startTime().unixTime()-modeltime)/3600.0)
                  self.VU.logMsg("Had to update MaxT at %d-hrs to match hourly T grids at %d points"%(fhr,changed))
                  self.createGrid(mutableModel,"MaxT","SCALAR",maxt,newtr)
      self.saveGrid(mutableModel,"MaxT")
      #
      #  Make sure Mint is as low as the min of all T grids
      #
      self.VU.logMsg("Checking for hourly T grids lower than MinT")
      gridInfoList=self.getGridInfo(mutableModel,"MinT","SFC",bigtr)
      for gridInfo in gridInfoList:
         newtr=gridInfo.gridTime()
         mintgrid=self.getGrids(mutableModel,"MinT","SFC",newtr,noDataError=0,cache=0)
         if mintgrid is not None:
            minoftgrid=self.getGrids(mutableModel,"T","SFC",newtr,mode="Min",noDataError=0,cache=0)
            if minoftgrid is not None:
               mint=minimum(minoftgrid,mintgrid)
               changed=count_nonzero(less(mint,mintgrid))
               if changed>0:
                  fhr=int((newtr.startTime().unixTime()-modeltime)/3600.0)
                  self.VU.logMsg("Had to update MinT at %d-hrs to match hourly T grids at %d points"%(fhr,changed))
                  self.createGrid(mutableModel,"MinT","SCALAR",mint,newtr)
      self.saveGrid(mutableModel,"MinT")
      #
      #  Make sure MaxRH is as high as the max of all RH grids
      #
      self.VU.logMsg("Checking for hourly RH grids higher than MaxRH")
      gridInfoList=self.getGridInfo(mutableModel,"MaxRH","SFC",bigtr)
      for gridInfo in gridInfoList:
         newtr=gridInfo.gridTime()
         maxrhgrid=self.getGrids(mutableModel,"MaxRH","SFC",newtr,noDataError=0,cache=0)
         if maxrhgrid is not None:
            maxofrhgrid=self.getGrids(mutableModel,"RH","SFC",newtr,mode="Max",noDataError=0,cache=0)
            if maxofrhgrid is not None:
               maxrh=maximum(maxofrhgrid,maxrhgrid)
               changed=count_nonzero(greater(maxrh,maxrhgrid))
               if changed>0:
                  fhr=int((newtr.startTime().unixTime()-modeltime)/3600.0)
                  self.VU.logMsg("Had to update MaxRH at %d-hrs to match hourly RH grids at %d points"%(fhr,changed))
                  self.createGrid(mutableModel,"MaxRH","SCALAR",maxrh,newtr)
      self.saveGrid(mutableModel,"MaxRH")
      #
      #  Make sure MinRH is as low as the min of all RH grids
      #
      self.VU.logMsg("Checking for hourly RH grids lower than MinRH")
      gridInfoList=self.getGridInfo(mutableModel,"MinRH","SFC",bigtr)
      for gridInfo in gridInfoList:
         newtr=gridInfo.gridTime()
         minrhgrid=self.getGrids(mutableModel,"MinRH","SFC",newtr,noDataError=0,cache=0)
         if minrhgrid is not None:
            minofrhgrid=self.getGrids(mutableModel,"RH","SFC",newtr,mode="Min",noDataError=0,cache=0)
            if minofrhgrid is not None:
               minrh=minimum(minofrhgrid,minrhgrid)
               changed=count_nonzero(less(minrh,minrhgrid))
               if changed>0:
                  fhr=int((newtr.startTime().unixTime()-modeltime)/3600.0)
                  self.VU.logMsg("Had to update MinRH at %d-hrs to match hourly RH grids at %d points"%(fhr,changed))
                  self.createGrid(mutableModel,"MinRH","SCALAR",minrh,newtr)
      self.saveGrid(mutableModel,"MinRH")
      #
      #  Create Td from T and RH
      #
      self.VU.logMsg("Making Td from T and RH")
      (self.parmUnits,self.parmPrecision,self.parmMinval,self.parmMaxval,
       self.parmColorTable,self.parmDisplayMinval,
       self.parmDisplayMaxval)=self.getParmInfo(mutableModel,"Td")
      Funits=0
      if self.parmUnits.find("F")>-1:
         Funits=1
      gridInfoList=self.getGridInfo(mutableModel,"T","SFC",bigtr)
      for gridInfo in gridInfoList:
         newtr=gridInfo.gridTime()
         tgrid=self.getGrids(mutableModel,"T","SFC",newtr,noDataError=0,cache=0)
         if tgrid is not None:
            rhgrid=self.getGrids(mutableModel,"RH","SFC",newtr,noDataError=0,cache=0)
            if rhgrid is not None:
               if Funits==1:
                  tc=(tgrid-32.0)*(5.0/9.0)
               else:
                  tc=tgrid
               rh=clip(rhgrid,0.5,99.5)/100.0
               x=(log(rh)/17.67)+(tc/(tc+243.5))
               tdc=(243.5*x)/(1-x)
               if Funits==1:
                  Td=(tdc*9.0/5.0)+32.0
               else:
                  Td=tdc
               Td=clip(Td,self.parmMinval,self.parmMaxval)
               self.createGrid(mutableModel,"Td","SCALAR",Td,newtr)
      self.saveGrid(mutableModel,"Td")
      #
      #  Calculate TdMrn/TdAft from MinT/MaxRH and MaxT/MinRH values
      #  Only if the TdMrn parameter exists for the mutable database
      #
      starttime=AbsTime.AbsTime(int(modeltime))
      parmInfoMrn = self.getParmInfo(mutableModel,"TdMrn")
      #if ((self.parmUnits!="NONE")or(self.parmPrecision!=-99)):
      if parmInfoMrn != None:
         self.VU.logMsg("Making TdMrn/TdAft grids")
         (self.parmUnits,self.parmPrecision,self.parmMinval,self.parmMaxval,
          self.parmColorTable,self.parmDisplayMinval,
          self.parmDisplayMaxval) = parmInfoMrn
         (runtr,gridTimes)=self.getGridTimes(mutableModel,"TdMrn","SFC",starttime,240)
         for newtr in gridTimes:
            maxRH=self.getGrids(mutableModel,"MaxRH","SFC",newtr,noDataError=0,cache=0)
            if maxRH is not None:
               minT=self.getGrids(mutableModel,"MinT","SFC",newtr,noDataError=0,cache=0)
               if minT is not None:
                  if Funits==1:
                     tc=(minT-32.0)*(5.0/9.0)
                  else:
                     tc=MinT
                  rh=clip(maxRH,0.5,99.5)/100.0
                  x=(log(rh)/17.67)+(tc/(tc+243.5))
                  tdc=(243.5*x)/(1-x)
                  if Funits==1:
                     Td=(tdc*9.0/5.0)+32.0
                  else:
                     Td=tdc
                  Td=clip(Td,self.parmMinval,self.parmMaxval)
                  self.createGrid(mutableModel,"TdMrn","SCALAR",Td,newtr)
         self.saveGrid(mutableModel,"TdMrn")
         (self.parmUnits,self.parmPrecision,self.parmMinval,self.parmMaxval,
          self.parmColorTable,self.parmDisplayMinval,
          self.parmDisplayMaxval)=self.getParmInfo(mutableModel,"TdAft")
         (runtr,gridTimes)=self.getGridTimes(mutableModel,"TdAft","SFC",starttime,240)
         for newtr in gridTimes:
            minRH=self.getGrids(mutableModel,"MinRH","SFC",newtr,noDataError=0,cache=0)
            if minRH is not None:
               maxT=self.getGrids(mutableModel,"MaxT","SFC",newtr,noDataError=0,cache=0)
               if maxT is not None:
                  if Funits==1:
                     tc=(maxT-32.0)*(5.0/9.0)
                  else:
                     tc=maxT
                  rh=clip(minRH,0.5,99.5)/100.0
                  x=(log(rh)/17.67)+(tc/(tc+243.5))
                  tdc=(243.5*x)/(1-x)
                  if Funits==1:
                     Td=(tdc*9.0/5.0)+32.0
                  else:
                     Td=tdc
                  Td=clip(Td,self.parmMinval,self.parmMaxval)
                  self.createGrid(mutableModel,"TdAft","SCALAR",Td,newtr)
         self.saveGrid(mutableModel,"TdAft")
      self.VU.logMsg("BOIVerifyBiasCorr Procedure Finished")
      return
   #=================================================================
   #  getLastUpdate(self,model,parmname,level,tr)
   #
   #  Gets the last update time for the grid for the model, parmname,
   #  level and timerange.  Should only call this for a timerange where
   #  only one grid is located. If multiple grids are found - gives the
   #  last update for the FIRST one in the list.  If no grids are found,
   #  it returns 0.
   #
   def getLastUpdate(self,model,parmname,level,tr):
      ret=-1
      try:
         historyList = self.getGridHistory(model, parmname, level, tr)
      except:
         return ret
      if not historyList:
         return ret
      for history in historyList[0]:
         if history[5] is not None:
            upd = history[5].unixTime()
         else:
            upd = None
         if upd>ret:
            ret=upd

      return ret
   #=================================================================
   #
   #  RH transforms
   #
   def Ztrans(self,ingrid):
      c=clip(ingrid,0.1,99.9)
      return 0.5*log(c/(100-c))
   def Rtrans(self,ingrid):
      v=exp(ingrid*2.0)
      return (100*v)/(1.0+v)
   #=================================================================
   #  getOldForecasts - Get list of matching forecast grids and
   #                    observed grids for the specified forecast
   #                    hour.
   #
   def getOldForecasts(self,caseInfo,fhrInfo,parm,model,obsparm,obsmodel,forecastHour,
                        minval,maxval):
      obsList=[]
      fcstList=[]
      #
      #
      #
      for key in caseInfo.keys():
         (freclist,oreclist)=caseInfo[key]
         #
         #  Have to have both forecast and observed grids
         #
         if ((len(freclist)<1)or(len(oreclist)<1)):
            continue
         #
         #  Use only ones that match the current forecastHour
         #
         fhr=fhrInfo[key]
         if fhr!=forecastHour:
            continue
         #
         #  Get average of non-bad forecast grids
         #
         numgrids=0
         fcstgrid=self.empty()
         for frec in freclist:
            testgrid=self.VU.readRecord(parm,model,frec)
            if self.badGrid(testgrid,minval,maxval):
               basetime=self.VU.fncBtime[frec]
               (byea,bmon,bday,bhou,bmin,bsec,bwda,byda,bdst)=time.gmtime(basetime)
               self.VU.logMsg("    %d-hr forecast from %4.4d/%2.2d/%2.2d %2.2dZ run appears bad - so skipped"%(fhr,
                              byea,bmon,bday,bhou),2)
               continue
            fcstgrid+=testgrid
            numgrids+=1
         if numgrids==0:
            continue
         fcstgrid/=float(numgrids)
         #
         #  Get average of non-bad observed grids
         #
         numgrids=0
         obsgrid=self.empty()
         for orec in oreclist:
            testgrid=self.VU.readRecord(obsparm,obsmodel,orec)
            if self.badGrid(testgrid,minval,maxval):
               starttime=self.VU.oncStime[orec]
               (syea,smon,sday,shou,smin,ssec,swda,syda,sdst)=time.gmtime(starttime)
               self.VU.logMsg("    Observed grid at %4.4d/%2.2d/%2.2d %2.2dZ appears bad - so skipped"%(syea,
                              smon,sday,shou),2)
               continue
            obsgrid+=testgrid
            numgrids+=1
         if numgrids==0:
            continue
         obsgrid/=float(numgrids)
         #
         #  Transform the ob/fcst if needed
         #
         if parm in self.transformParms:
            obsgrid=self.Ztrans(obsgrid)
            fcstgrid=self.Ztrans(fcstgrid)
         #
         #  Add grid to lists of grids
         #
         obsList.append(obsgrid)
         fcstList.append(fcstgrid)
      #
      #  Return lists of forecast/observed grids
      #
      return(obsList,fcstList)
   #=================================================================
   #  getRegression
   #
   #  get linear regression predictors of non-average forecast errors
   #  from forecast anomalies (difference of forecast from average
   #  forecast) from the pairList forecasts
   #
   def getRegression(self,avgFcst,avgErr,obsGrids,fcstGrids):
      #
      #  The independent variable is forecast anomaly (forecast-avgFcst)
      #
      ind=subtract(fcstGrids,avgFcst)
      #
      #  The dependent variable is error anomaly (error-avgErr)
      #
      dep=subtract(fcstGrids-obsGrids,avgErr)
      #
      #  Get covariance, and square of dependent and independent variables
      #  across the cases
      #
      covsum=add.reduce(dep*ind)
      depsqr=add.reduce(dep*dep)
      indsqr=add.reduce(ind*ind)
      #
      #  Calculate standard deviation of indpendent variables
      #
      num=obsGrids.shape[0]
      stdind=(indsqr/num)**0.5 # standard deviation of independent variable
      stddep=(depsqr/num)**0.5 # standard deviation of dependent variable
      stdind[less(stdind,0.001)] = 0.001
      stddep[less(stddep,0.001)] = 0.001
      covariance=covsum/num    # covariance
      correlation=covariance/(stdind*stddep)  # correlation coefficient
      return(correlation,stdind,stddep)
   #==================================================================
   #
   #  badGrid - very crude way to flag bad grids from calculations.
   #            Just looks for grids with virtually no variance.
   #            (this won't work for precip that typically does have
   #             no variance on days with no precip)
   #            returns 1 if bad, 0 if OK.
   #
   def badGrid(self,grid,minvalue,maxvalue):
      gridShape = self.getGridShape()
      numpts=gridShape[0]*gridShape[1]
      avg=add.reduce(add.reduce(grid))/float(numpts)
      sqr=add.reduce(add.reduce(grid*grid))/float(numpts)
      var=sqr-(avg*avg)
      if var<((maxvalue-minvalue)/1000.0):
         return 1
      return 0
   #==================================================================
   #
   #
   def getParmInfo(self,mutableModel,parm):
      units="NONE"
      precision=-99
      minval=0
      maxval=100
      colorTable=""
      displayMinval=0
      displayMaxval=100
      parm=self.getParm(mutableModel,parm,"SFC")
      if parm is not None:
         parmInfo = parm.getGridInfo()
         units=parmInfo.getUnitString()
         precision=parmInfo.getPrecision()
         minval=parmInfo.getMinValue()
         maxval=parmInfo.getMaxValue()
         if DSPATTR_PORT:
            ctInfo=parmInfo.getDisplayAttributes().colorTable()
            if ctInfo is not None:
               colorTable=ctInfo.name()
               displayMinval=ctInfo.minLimit()
               displayMaxval=ctInfo.maxLimit()
         return(units,precision,minval,maxval,colorTable,
             displayMinval,displayMaxval)
   #-----------------------------------------------------------------
   #
   #  Save grid for the specified outdb and parmname
   #
   def saveGrid(self,outdb,parmname):
      p=self.getParm(outdb,parmname,"SFC")
      if p is not None:
         p.saveParameter(1)
      return
#
#  debug stuff for memory usage
#
_proc_status="/proc/%d/status"%os.getpid()
_scale={'kB':1024.0,'mB':1024.0*1024.0,
        'KB':1024.0,'MB':1024.0*1024.0}
def _VmB(VmKey):
   try:
      t=open(_proc_status)
      v=t.read()
      t.close()
   except IOError:
      return 0.0
   i=v.index(VmKey)
   v=v[i:].split(None,3)
   if len(v)<3:
      return 0.0
   return float(v[1])*_scale[v[2]]
def memory():
   return _VmB('VmSize:')
def resident():
   return _VmB('VmRSS:')

ext: py

reason: from numpy


text:
"""Support for globally optimizing tensor networks using tensorflow in either
interactive (eager) mode or compiled (session) mode. For computationally heavy
optimizations (complex contractions) the two modes offer similar performance,
but for small to medium contractions graph mode is significantly faster.
"""

import time
import functools
from collections import namedtuple

import tqdm
import numpy as np

from .tensor_core import TensorNetwork

LazyComplexTF = namedtuple('LazyComplexTF', ['shape', 'real', 'imag'])
"""We need to lazily represent complex tensorflow arrays so that combining the
variable real and complex parts takes place in the gradient tape always when
executing eagerly.
"""

_TENSORFLOW = None
_TENSORFLOW_EAGER = None

def get_tensorflow(eager=None):
    """Get tensorflow and initialize eager mode if tensorflow has not been
    imported yet. If it has, check whether eager mode has been enabled.
    """
    global _TENSORFLOW
    global _TENSORFLOW_EAGER

    if _TENSORFLOW is None:
        # first check if we are importing tensorflow for the first time
        import sys

        try:
            # leave eager decision to whenever tensorflow was imported
            _TENSORFLOW = sys.modules['tensorflow']
            _TENSORFLOW_EAGER = _TENSORFLOW.executing_eagerly()

        except KeyError:
            import tensorflow as tf

            # default to enabling the more intuitive eager mode
            if eager is None:
                eager = True

            if eager:
                tf.compat.v1.enable_eager_execution()

            _TENSORFLOW = tf
            _TENSORFLOW_EAGER = eager

    return _TENSORFLOW

def executing_eagerly():
    global _TENSORFLOW_EAGER
    return _TENSORFLOW_EAGER

def variable_tn(tn):
    """Convert a tensor network's arrays to tensorflow Variables.
    """
    tf = get_tensorflow()
    tf_tn = tn.copy()
    tf_tn.apply_to_arrays(tf.Variable)
    return tf_tn

def constant_tn(tn):
    """Convert a tensor network's arrays to tensorflow constants.
    """
    tf = get_tensorflow()
    tf_tn = tn.copy()
    tf_tn.apply_to_arrays(tf.convert_to_tensor)
    return tf_tn

def parse_network_to_tf(tn, constant_tags):
    tf = get_tensorflow()
    eager = executing_eagerly()

    tn_tf = tn.copy()
    variables = []
    iscomplex = False

    for t in tn_tf:

        # check if tensor has any of the constant tags
        if t.tags & constant_tags:
            t.modify(data=tf.convert_to_tensor(value=t.data))

        # treat re and im parts as separate variables
        elif issubclass(t.dtype.type, np.complexfloating):
            tf_re, tf_im = tf.Variable(t.data.real), tf.Variable(t.data.imag)
            variables.extend((tf_re, tf_im))

            # if eager need to delay complex combination until gradient tape
            if eager:
                tf_data = LazyComplexTF(t.data.shape, tf_re, tf_im)
            else:
                tf_data = tf.complex(tf_re, tf_im)

            t.modify(data=tf_data)
            iscomplex = True

        else:
            tf_data = tf.Variable(t.data)
            variables.append(tf_data)
            t.modify(data=tf_data)

    return tn_tf, variables, iscomplex

def evaluate_lazy_complex(tn_tf):
    """Can call this in a gradient tape context to merge real and imaginary
    parts so far lazily represented.
    """
    tf = get_tensorflow()

    tn_eval = tn_tf.copy()

    for t in tn_eval:
        if isinstance(t.data, LazyComplexTF):
            t.modify(data=tf.complex(t.data.real, t.data.imag))

    return tn_eval

def init_uninit_vars(sess):
    """ Initialize all other trainable variables, i.e. those which are
    uninitialized. Allows ``optimize`` to be interatively stopped and started.
    """
    tf = get_tensorflow()

    # names of variables
    uninit_vars = sess.run(tf.compat.v1.report_uninitialized_variables())
    vars_list = {v.decode('utf-8') for v in uninit_vars}

    # variables themselves
    uninit_vars_tf = [v for v in tf.compat.v1.global_variables()
                      if v.name.split(':')[0] in vars_list]

    sess.run(tf.compat.v1.variables_initializer(var_list=uninit_vars_tf))

class TNOptimizer:
    """Optimize a tensor network's arrays using tensorflow.

    Parameters
    ----------
    tn : TensorNetwork
        The tensor network to optimize, every tensor in this will be treated
        as a variable.
    loss_fn : callable
        The target loss function to minimize. Should accept ``tn`` (backed by
        tensorflow tensors) or the output of ``norm_fn``.
    norm_fn : callable, optional
        A function to normalize ``tn`` before being passed to ``loss_fn`` and
        also to call on the final, optimized tensor network. This can be used
        to enforce constraints such as normalization or unitarity.
    loss_constants : dict_like, optional
        Extra constant arguments to supply to ``loss_fn`` and be converted to
        tensorflow constant tensors. Can be individual arrays or tensor
        networks.
    loss_kwargs : dict_like, optional
        Other kwargs to supply to ``loss_fn`` that are not arrays or tensor
        networks.
    constant_tags : sequence of str, optional
        Treat any tensors *within* ``tn`` with these tags as constant.
    optimizer : str, optional
        Which optimizer to use, default: ``'AdamOptimizer'``.
        This should be an optimizer that can be found in the
        ``tensorflow.train`` submodule. The exception is if ``scipy`` is
        supplied the ``tf.contrib.opt.ScipyOptimizerInterface`` will be used,
        the options of which should be specified in the ``optimize`` call.
        This is mainly useful for the default ``'L-BFGS-B'`` algorithm.
    learning_rate : float, optional
        The learning rate to apply to use the optimizer with. You can
        dynamically change this between ``optimize`` calls.
    learning_decay_steps : int, optional
        How many steps to decay the learning rate over.
    learning_decay_rate : float, optional
        How much to decay the learning rate over ``learning_decay_steps``.
    loss_target : float, optional
        If supplied, stop optimizing once this loss is reached.
    progbar : bool, optional
        Whether to show live progress of the optimization.

    Examples
    --------

    .. code:: python3

        import quimb.tensor as qtn
        from quimb.tensor.optimize_tensorflow import TNOptimizer

        # can leave this out, and quimb will invoke tensorflow in
        # eager mode, sacrificing some efficiency
        import tensorflow as tf
        sess = tf.InteractiveSession()

    Variational compression of cyclic MPS, halving its bond dimension:

    .. code:: python3

        # the target state we want to compress
        targ = qtn.MPS_rand_state(n=20, bond_dim=16,
                                  cyclic=True, dtype='complex64')

        # our initial guess for the compressed state
        psi0 = qtn.MPS_rand_state(n=20, bond_dim=8,
                                  cyclic=True, dtype='complex64')

    .. code:: python3

        >>> targ.show()
         16 16 16 16 16 16 16 16 16 16 16 16 16 16 16 16 16 16 16 16 16
        +--o--o--o--o--o--o--o--o--o--o--o--o--o--o--o--o--o--o--o--o--+
           |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |

        >>> psi0.show()
         8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8
        +-o-o-o-o-o-o-o-o-o-o-o-o-o-o-o-o-o-o-o-o-+
          | | | | | | | | | | | | | | | | | | | |

        >>> abs(psi0.H @ targ)
        0.0012382069

    .. code:: python3

        # the thing (scalar quantity) we want to *minimize*
        def abs_overlap(psi, target):
            overlap = (psi.H & target) ^ all
            return -abs(overlap)

        # also need to define function that keeps us in the normalized manifold
        def psi_normalize(psi):
            nfactor = (psi.H & psi) ^ all
            return psi / nfactor**0.5

    Now we can set-up the optimizer and run it:

    .. code:: python3

        tnopt = TNOptimizer(
            tn=psi0,
            loss_fn=abs_overlap,
            loss_constants={'target': targ},
            norm_fn=psi_normalize,
            optimizer='Adam',
        )

    .. code:: python3

        >>> psif = tnopt.optimize(400)
        -0.7474728226661682: 100%|███████████| 400/400 [00:07<00:00, 51.74it/s]

        >>> abs(psif.H @ targ)
        0.747473

    You could alter the learning rate and keep running this if satisfactory
    results are not achieved. Note we don't here expect this to reach 1.0 as
    for random MPS all the singular values are likely significant and so
    compressing always involves some loss of fidelity.
    """

    def __init__(
        self,
        tn,
        loss_fn,
        norm_fn=None,
        loss_constants=None,
        loss_kwargs=None,
        constant_tags=None,
        optimizer='AdamOptimizer',
        learning_rate=0.1,
        learning_decay_steps=100,
        learning_decay_rate=0.5,
        loss_target=None,
        progbar=True
    ):
        tf = get_tensorflow()

        self.constant_tags = (set() if constant_tags is None
                              else set(constant_tags))

        # make tensorflow version of network and gather variables etc.
        self.tn_opt, self.variables, self.iscomplex = parse_network_to_tf(
            tn, self.constant_tags
        )

        self.progbar = progbar
        self.loss_target = loss_target

        # set dynamic learning rate
        self.global_step = tf.compat.v1.train.get_or_create_global_step()
        self.learning_rate = tf.compat.v1.train.exponential_decay(
            learning_rate=learning_rate, global_step=self.global_step,
            decay_steps=learning_decay_steps, decay_rate=learning_decay_rate,
        )

        # use identity if no nomalization required
        if norm_fn is None:
            def norm_fn(x):
                return x

        self.norm_fn = norm_fn
        self.loss_constants = {}
        if loss_constants is not None:
            for k, v in loss_constants.items():
                # check if tensor network supplied
                if isinstance(v, TensorNetwork):
                    # convert it to constant tensorflow TN
                    self.loss_constants[k] = constant_tn(v)
                else:
                    self.loss_constants[k] = tf.convert_to_tensor(value=v)
        self.loss_kwargs = {} if loss_kwargs is None else dict(loss_kwargs)
        self.loss_fn = functools.partial(
            loss_fn, **self.loss_constants, **self.loss_kwargs
        )
        self.loss = None
        self._n = 0

        if optimizer == 'scipy':
            if executing_eagerly():
                raise ValueError("The tensorflow ``'scipy'`` interface is not "
                                 "available when executing eagerly.")

            self.optimizer = 'scipy'
            return

        elif isinstance(optimizer, str):
            if 'Optimizer' not in optimizer:
                optimizer += 'Optimizer'
            self.optimizer = getattr(
                tf.compat.v1.train, optimizer)(self.learning_rate)
        else:
            self.optimizer = optimizer

        if not executing_eagerly():
            # generate the loss and optimizer computational graphs
            self.loss_op = self.loss_fn(self.norm_fn(self.tn_opt))
            self.train_op = self.optimizer.minimize(
                self.loss_op, global_step=self.global_step)

    @property
    def nevals(self):
        """The number of gradient evaluations.
        """
        return self._n

    def _maybe_start_timer(self, max_time):
        if max_time is None:
            return
        else:
            self._time_start = time.time()

    def _time_should_stop(self, max_time):
        if max_time is None:
            return False
        else:
            return (time.time() - self._time_start) > max_time

    def _get_tn_opt_numpy_eager(self):
        """Get normalized, numpy version of optimized tensor network, in eager
        mode.
        """
        if self.iscomplex:
            tn = evaluate_lazy_complex(self.tn_opt)
        else:
            tn = self.tn_opt

        tn_opt_numpy = self.norm_fn(tn)
        tn_opt_numpy.apply_to_arrays(lambda x: x.numpy())
        return tn_opt_numpy

    def _get_tn_opt_numpy_graph(self, sess):
        """Get normalized, numpy version of optimized tensor network.
        """
        tn_opt_numpy = self.norm_fn(self.tn_opt)

        # evaluate all arrays as once for performance reasons
        np_arrays = sess.run([t.data for t in tn_opt_numpy])
        for t, data in zip(tn_opt_numpy, np_arrays):
            t.modify(data=data)

        return tn_opt_numpy

    def _optimize_scipy(self, max_steps, method='L-BFGS-B',
                        sess=None, max_time=None, **kwargs):
        tf = get_tensorflow()

        loss_op = self.loss_fn(self.norm_fn(self.tn_opt))

        kwargs['maxfun'] = max_steps - 1
        if sess is None:
            sess = tf.compat.v1.get_default_session()
        init_uninit_vars(sess)

        pbar = tqdm.tqdm(total=max_steps, disable=not self.progbar)
        self._maybe_start_timer(max_time)
        timed_out = False

        def step_callback(*_, **__):
            if self._time_should_stop(max_time):
                raise TimeoutError

        def loss_callback(loss_val):
            self._n += 1
            self.loss = loss_val
            pbar.set_description(f"{loss_val}")
            pbar.update()
            return self._time_should_stop(max_time)

        try:
            optimizer = tf.contrib.opt.ScipyOptimizerInterface(
                loss_op, options=kwargs, method=method,
            )
            optimizer.minimize(sess, fetches=[loss_op],
                               loss_callback=loss_callback,
                               step_callback=step_callback)
        except TimeoutError:
            timed_out = True
        finally:
            pbar.close()

            if timed_out:
                raise TimeoutError("The scipy optimizer currently has no way "
                                   "of extracting the current tensor values "
                                   "mid execution.")

            return self._get_tn_opt_numpy_graph(sess)

    def _grad(self):
        tf = get_tensorflow()
        # this is the function that computes the loss and tracks gradient
        with tf.GradientTape() as tape:
            if self.iscomplex:
                tn = evaluate_lazy_complex(self.tn_opt)
            else:
                tn = self.tn_opt
            self.loss = self.loss_fn(self.norm_fn(tn))
        return tape.gradient(self.loss, self.variables)

    def _optimize_eager(self, max_steps, max_time=None):
        # perform the optimization with live progress
        pbar = tqdm.tqdm(total=max_steps, disable=not self.progbar)
        self._maybe_start_timer(max_time)
        try:
            for _ in range(max_steps):
                # compute gradient and display loss value
                grads = self._grad()
                pbar.set_description(f"{self.loss}")

                # performing an optimization step
                self.optimizer.apply_gradients(
                    zip(grads, self.variables),
                    global_step=self.global_step,
                )
                pbar.update()
                self._n += 1

                # check stopping criteria
                if (self.loss_target is not None):
                    if self.loss < self.loss_target:
                        break
                if self._time_should_stop(max_time):
                    break

        finally:
            pbar.close()
            return self._get_tn_opt_numpy_eager()

    def _optimize_graph(self, max_steps, max_time=None, sess=None, ):
        if sess is None:
            tf = get_tensorflow()
            sess = tf.compat.v1.get_default_session()

        init_uninit_vars(sess)

        # perform the optimization with live progress
        pbar = tqdm.tqdm(total=max_steps, disable=not self.progbar)
        self._maybe_start_timer(max_time)
        try:
            for _ in range(max_steps):
                _, self.loss = sess.run([self.train_op, self.loss_op])
                pbar.set_description(f"{self.loss}")
                pbar.update()
                self._n += 1

                # check stopping criteria
                if (self.loss_target is not None):
                    if self.loss < self.loss_target:
                        break
                if self._time_should_stop(max_time):
                    break
        except Exception as e:
            print(e)
        finally:
            pbar.close()
            return self._get_tn_opt_numpy_graph(sess)

    def optimize(self, max_steps, max_time=None, **kwargs):
        """
        Returns
        -------
        tn_opt : TensorNetwork
            The optimized tensor network (with arrays converted back to numpy).
        """
        if self.optimizer == 'scipy':
            kwargs.setdefault('gtol', 1e-12)
            return self._optimize_scipy(max_steps, max_time=max_time, **kwargs)
        elif executing_eagerly():
            return self._optimize_eager(max_steps, max_time=max_time, **kwargs)
        else:
            return self._optimize_graph(max_steps, max_time=max_time, **kwargs)

max_stars_repo_name: vonDonnerstein/quimb

max_issues_repo_name: vonDonnerstein/quimb

max_forks_repo_name: vonDonnerstein/quimb

reason: import numpy


text:
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import random
from pathlib import Path

import numpy as np
from paddle.io import BatchSampler
from paddle.io import Dataset

from paddlespeech.t2s.exps.ge2e.random_cycle import random_cycle

class MultiSpeakerMelDataset(Dataset):
    """A 2 layer directory thatn contains mel spectrograms in *.npy format.
    An Example file structure tree is shown below. We prefer to preprocess
    raw datasets and organized them like this.

    dataset_root/
      speaker1/
        utterance1.npy
        utterance2.npy
        utterance3.npy
      speaker2/
        utterance1.npy
        utterance2.npy
        utterance3.npy
    """

    def __init__(self, dataset_root: Path):
        self.root = Path(dataset_root).expanduser()
        speaker_dirs = [f for f in self.root.glob("*") if f.is_dir()]

        speaker_utterances = {
            speaker_dir: list(speaker_dir.glob("*.npy"))
            for speaker_dir in speaker_dirs
        }

        self.speaker_dirs = speaker_dirs
        self.speaker_to_utterances = speaker_utterances

        # meta data
        self.num_speakers = len(self.speaker_dirs)
        self.num_utterances = np.sum(
            len(utterances)
            for speaker, utterances in self.speaker_to_utterances.items())

    def get_example_by_index(self, speaker_index, utterance_index):
        speaker_dir = self.speaker_dirs[speaker_index]
        fpath = self.speaker_to_utterances[speaker_dir][utterance_index]
        return self[fpath]

    def __getitem__(self, fpath):
        return np.load(fpath)

    def __len__(self):
        return int(self.num_utterances)

class MultiSpeakerSampler(BatchSampler):
    """A multi-stratal sampler designed for speaker verification task.
    First, N speakers from all speakers are sampled randomly. Then, for each
    speaker, randomly sample M utterances from their corresponding utterances.
    """

    def __init__(self,
                 dataset: MultiSpeakerMelDataset,
                 speakers_per_batch: int,
                 utterances_per_speaker: int):
        self._speakers = list(dataset.speaker_dirs)
        self._speaker_to_utterances = dataset.speaker_to_utterances

        self.speakers_per_batch = speakers_per_batch
        self.utterances_per_speaker = utterances_per_speaker

    def __iter__(self):
        # yield list of Paths
        speaker_generator = iter(random_cycle(self._speakers))
        speaker_utterances_generator = {
            s: iter(random_cycle(us))
            for s, us in self._speaker_to_utterances.items()
        }

        while True:
            speakers = []
            for _ in range(self.speakers_per_batch):
                speakers.append(next(speaker_generator))

            utterances = []
            for s in speakers:
                us = speaker_utterances_generator[s]
                for _ in range(self.utterances_per_speaker):
                    utterances.append(next(us))
            yield utterances

class RandomClip(object):
    def __init__(self, frames):
        self.frames = frames

    def __call__(self, spec):
        # spec [T, C]
        T = spec.shape[0]
        start = random.randint(0, T - self.frames)
        return spec[start:start + self.frames, :]

class Collate(object):
    def __init__(self, num_frames):
        self.random_crop = RandomClip(num_frames)

    def __call__(self, examples):
        frame_clips = [self.random_crop(mel) for mel in examples]
        batced_clips = np.stack(frame_clips)
        return batced_clips

if __name__ == "__main__":
    mydataset = MultiSpeakerMelDataset(
        Path("/home/chenfeiyu/datasets/SV2TTS/encoder"))
    print(mydataset.get_example_by_index(0, 10))

max_stars_repo_name: gongel/DeepSpeech

max_issues_repo_name: gongel/DeepSpeech

max_forks_repo_name: gongel/DeepSpeech


text:
import os
import numpy as np
from ase.io import read
import matplotlib.pyplot as plt

with open("./hessian") as out_file:
    out_file.readline()

    lines = out_file.readlines()

    prev_atom_idx = 0
    forces = []
    tmp_forces = []
    for idx, line in enumerate(lines):
        if "$end" in line:
            forces.append(np.array(tmp_forces))
            tmp_forces = []
            break

        line = line.split()
        new_atom_idx = int(line[0])

        if prev_atom_idx == new_atom_idx or new_atom_idx == 1:
            for val in line[2:]:
                tmp_forces.append(float(val))
        else:
            forces.append(np.array(tmp_forces))

            tmp_forces = []
            for val in line[2:]:
                tmp_forces.append(float(val))
        prev_atom_idx = new_atom_idx
    forces = np.array(forces)

forces = forces[0::3, 0::3]
plt.imshow(forces, origin="lower")
plt.show()
np.save("./turbomole_force_constants.npy", forces)


text:
import numpy as np
from psychopy import visual, core, event
from psychopy.visual.circle import Circle
from toon.input.mpdevice import MpDevice
from toon.input.mouse import Mouse
from toon.anim.easing import linear
from toon.anim import Timeline, Track, Player

class AnimCircle(Player, Circle):
    def __init__(self, *args, **kwargs):
        Player.__init__(self)
        Circle.__init__(self, *args, **kwargs)

win = visual.Window(fullscr=True, units='pix', screen=1)

timeline = Timeline()
timeline.start()
psy_mouse = event.Mouse()
toon_mouse = MpDevice(Mouse())

target_cir = AnimCircle(win, size=100, fillColor='white', lineColor=None)
toon_cir = visual.Circle(win, size=20, fillColor='blue', opacity=0.5, lineColor=None)
toon_cir.pos = (-win.size[0]/2, win.size[1]/2)

times = np.arange(0, 10, 1/60)
keyframes = np.vstack((times, np.sin(times) * 400)).T
keyframes2 = np.vstack((times, np.cos(times) * 400)).T
x_track = Track(keyframes, easing=linear)
y_track = Track(keyframes2, easing=linear)

def change_x(val, obj):
    obj.pos = [val, obj.pos[1]]

def change_y(val, obj):
    obj.pos = [obj.pos[0], val]

target_cir.add(x_track, change_x)
target_cir.add(y_track, change_y)
timeline.next_frame()
target_cir.start(timeline.frame_time)
with toon_mouse:  # alternatively, use toon_mouse.start(), toon_mouse.stop()
    while target_cir.is_playing and not event.getKeys(['esc', 'escape']):
        win.callOnFlip(timeline.next_frame)
        target_cir.advance(timeline.frame_time)
        clicks, pos, scroll = toon_mouse.read()
        if pos is not None:
            pos *= [1, -1]
            for i in pos:
                toon_cir.pos += i
        if toon_cir.overlaps(target_cir):
            target_cir.fillColor = 'lightgray'
        else:
            target_cir.fillColor = 'white'
        target_cir.draw()
        toon_cir.draw()
        win.flip()

win.close()
core.quit()


text:
from setuptools import setup, find_packages
import numpy as np

setup(
    name='ppo-ewma',
    packages=find_packages(),
    version='0.0.1',
)


text:
# -*- coding: utf-8 -*-
from __future__ import division
import numpy as np
import torch

Transition_dtype = np.dtype([('timestep', np.int32), ('state', np.uint8, (3, 64, 64)), ('action', np.int32), ('reward', np.float32), ('nonterminal', np.bool_)])
blank_trans = (0, np.zeros((3, 64, 64), dtype=np.uint8), 0, 0.0, False)

# Segment tree data structure where parent node values are sum/max of children node values
class SegmentTree():
  def __init__(self, size):
    self.index = 0
    self.size = size
    self.full = False  # Used to track actual capacity
    self.tree_start = 2**(size-1).bit_length()-1  # Put all used node leaves on last tree level
    self.sum_tree = np.zeros((self.tree_start + self.size,), dtype=np.float32)
    self.data = np.array([blank_trans] * size, dtype=Transition_dtype)  # Build structured array
    self.max = 1  # Initial max value to return (1 = 1^ω)

  # Updates nodes values from current tree
  def _update_nodes(self, indices):
    children_indices = indices * 2 + np.expand_dims([1, 2], axis=1)
    self.sum_tree[indices] = np.sum(self.sum_tree[children_indices], axis=0)

  # Propagates changes up tree given tree indices
  def _propagate(self, indices):
    parents = (indices - 1) // 2
    unique_parents = np.unique(parents)
    self._update_nodes(unique_parents)
    if parents[0] != 0:
      self._propagate(parents)

  # Propagates single value up tree given a tree index for efficiency
  def _propagate_index(self, index):
    parent = (index - 1) // 2
    left, right = 2 * parent + 1, 2 * parent + 2
    self.sum_tree[parent] = self.sum_tree[left] + self.sum_tree[right]
    if parent != 0:
      self._propagate_index(parent)

  # Updates values given tree indices
  def update(self, indices, values):
    self.sum_tree[indices] = values  # Set new values
    self._propagate(indices)  # Propagate values
    current_max_value = np.max(values)
    self.max = max(current_max_value, self.max)

  # Updates single value given a tree index for efficiency
  def _update_index(self, index, value):
    self.sum_tree[index] = value  # Set new value
    self._propagate_index(index)  # Propagate value
    self.max = max(value, self.max)

  def append(self, data, value):
    self.data[self.index] = data  # Store data in underlying data structure
    self._update_index(self.index + self.tree_start, value)  # Update tree
    self.index = (self.index + 1) % self.size  # Update index
    self.full = self.full or self.index == 0  # Save when capacity reached
    self.max = max(value, self.max)

  # Searches for the location of values in sum tree
  def _retrieve(self, indices, values):
    children_indices = (indices * 2 + np.expand_dims([1, 2], axis=1)) # Make matrix of children indices
    # If indices correspond to leaf nodes, return them
    if children_indices[0, 0] >= self.sum_tree.shape[0]:
      return indices
    # If children indices correspond to leaf nodes, bound rare outliers in case total slightly overshoots
    elif children_indices[0, 0] >= self.tree_start:
      children_indices = np.minimum(children_indices, self.sum_tree.shape[0] - 1)
    left_children_values = self.sum_tree[children_indices[0]]
    successor_choices = np.greater(values, left_children_values).astype(np.int32)  # Classify which values are in left or right branches
    successor_indices = children_indices[successor_choices, np.arange(indices.size)] # Use classification to index into the indices matrix
    successor_values = values - successor_choices * left_children_values  # Subtract the left branch values when searching in the right branch
    return self._retrieve(successor_indices, successor_values)

  # Searches for values in sum tree and returns values, data indices and tree indices
  def find(self, values):
    indices = self._retrieve(np.zeros(values.shape, dtype=np.int32), values)
    data_index = indices - self.tree_start
    return (self.sum_tree[indices], data_index, indices)  # Return values, data indices, tree indices

  # Returns data given a data index
  def get(self, data_index):
    return self.data[data_index % self.size]

  def total(self):
    return self.sum_tree[0]

class ReplayMemory():
  def __init__(self, args, capacity):
    self.device = args.device
    self.capacity = capacity
    self.history = args.history_length
    self.discount = args.discount
    self.n = args.multi_step
    self.priority_weight = args.priority_weight  # Initial importance sampling weight β, annealed to 1 over course of training
    self.priority_exponent = args.priority_exponent
    self.t = 0  # Internal episode timestep counter
    self.n_step_scaling = torch.tensor([self.discount ** i for i in range(self.n)], dtype=torch.float32, device=self.device)  # Discount-scaling vector for n-step returns
    self.transitions = SegmentTree(capacity)  # Store transitions in a wrap-around cyclic buffer within a sum tree for querying priorities

  # Adds state and action at time t, reward and terminal at time t + 1
  def append(self, state, action, reward, terminal):
    state = state[-1].mul(255).to(dtype=torch.uint8, device=torch.device('cpu'))  # Only store last frame and discretise to save memory
    self.transitions.append((self.t, state, action, reward, not terminal), self.transitions.max)  # Store new transition with maximum priority
    self.t = 0 if terminal else self.t + 1  # Start new episodes with t = 0

  # Returns the transitions with blank states where appropriate
  def _get_transitions(self, idxs):
    transition_idxs = np.arange(-self.history + 1, self.n + 1) + np.expand_dims(idxs, axis=1)
    transitions = self.transitions.get(transition_idxs)
    transitions_firsts = transitions['timestep'] == 0
    blank_mask = np.zeros_like(transitions_firsts, dtype=np.bool_)
    for t in range(self.history - 2, -1, -1):  # e.g. 2 1 0
      blank_mask[:, t] = np.logical_or(blank_mask[:, t + 1], transitions_firsts[:, t + 1]) # True if future frame has timestep 0
    for t in range(self.history, self.history + self.n):  # e.g. 4 5 6
      blank_mask[:, t] = np.logical_or(blank_mask[:, t - 1], transitions_firsts[:, t]) # True if current or past frame has timestep 0
    transitions[blank_mask] = blank_trans
    return transitions

  # Returns a valid sample from each segment
  def _get_samples_from_segments(self, batch_size, p_total):
    segment_length = p_total / batch_size  # Batch size number of segments, based on sum over all probabilities
    segment_starts = np.arange(batch_size) * segment_length
    valid = False
    while not valid:
      samples = np.random.uniform(0.0, segment_length, [batch_size]) + segment_starts  # Uniformly sample from within all segments
      probs, idxs, tree_idxs = self.transitions.find(samples)  # Retrieve samples from tree with un-normalised probability
      if np.all((self.transitions.index - idxs) % self.capacity > self.n) and np.all((idxs - self.transitions.index) % self.capacity >= self.history) and np.all(probs != 0):
        valid = True  # Note that conditions are valid but extra conservative around buffer index 0
    # Retrieve all required transition data (from t - h to t + n)
    transitions = self._get_transitions(idxs)
    # Create un-discretised states and nth next states
    all_states = transitions['state']
    states = torch.tensor(all_states[:, :self.history], device=self.device, dtype=torch.float32).div_(255)
    next_states = torch.tensor(all_states[:, self.n:self.n + self.history], device=self.device, dtype=torch.float32).div_(255)
    # Discrete actions to be used as index
    actions = torch.tensor(np.copy(transitions['action'][:, self.history - 1]), dtype=torch.int64, device=self.device)
    # Calculate truncated n-step discounted returns R^n = Σ_k=0->n-1 (γ^k)R_t+k+1 (note that invalid nth next states have reward 0)
    rewards = torch.tensor(np.copy(transitions['reward'][:, self.history - 1:-1]), dtype=torch.float32, device=self.device)
    R = torch.matmul(rewards, self.n_step_scaling)
    # Mask for non-terminal nth next states
    nonterminals = torch.tensor(np.expand_dims(transitions['nonterminal'][:, self.history + self.n - 1], axis=1), dtype=torch.float32, device=self.device)
    return probs, idxs, tree_idxs, states, actions, R, next_states, nonterminals

  def sample(self, batch_size):
    p_total = self.transitions.total()  # Retrieve sum of all priorities (used to create a normalised probability distribution)
    probs, idxs, tree_idxs, states, actions, returns, next_states, nonterminals = self._get_samples_from_segments(batch_size, p_total)  # Get batch of valid samples
    probs = probs / p_total  # Calculate normalised probabilities
    capacity = self.capacity if self.transitions.full else self.transitions.index
    weights = (capacity * probs) ** -self.priority_weight  # Compute importance-sampling weights w
    weights = torch.tensor(weights / weights.max(), dtype=torch.float32, device=self.device)  # Normalise by max importance-sampling weight from batch
    return tree_idxs, states, actions, returns, next_states, nonterminals, weights

  def update_priorities(self, idxs, priorities):
    priorities = np.power(priorities, self.priority_exponent)
    self.transitions.update(idxs, priorities)

  # Set up internal state for iterator
  def __iter__(self):
    self.current_idx = 0
    return self

  # Return valid states for validation
  def __next__(self):
    if self.current_idx == self.capacity:
      raise StopIteration
    transitions = self.transitions.data[np.arange(self.current_idx - self.history + 1, self.current_idx + 1)]
    transitions_firsts = transitions['timestep'] == 0
    blank_mask = np.zeros_like(transitions_firsts, dtype=np.bool_)
    for t in reversed(range(self.history - 1)):
      blank_mask[t] = np.logical_or(blank_mask[t + 1], transitions_firsts[t + 1]) # If future frame has timestep 0
    transitions[blank_mask] = blank_trans
    state = torch.tensor(transitions['state'], dtype=torch.float32, device=self.device).div_(255)  # Agent will turn into batch
    self.current_idx += 1
    return state

  next = __next__  # Alias __next__ for Python 2 compatibility

max_stars_repo_name: TemplierPaul/Rainbow

max_issues_repo_name: TemplierPaul/Rainbow

max_forks_repo_name: TemplierPaul/Rainbow


text:
import numpy as np
from gensim.models import KeyedVectors

import os
import json
import subprocess

FASTTEXT_CRAWL_NAME = 'fasttext_crawl.2M.300d'
FASTTEXT_WIKI_NAME = 'fasttext_wiki.1M.300d'
GLOVE_NAME = 'glove.6B.50d'
SENNA_NAME = 'senna.50d'

class WordEmbeddingLoader(object):
    """
    A loader for pre-trained word embedding
    """

    def __init__(self, embedding_name):
        self.embedding_name = embedding_name
        if 'glove' == embedding_name:
            self.word_dim = 50
            self.path_word = 'pretrain/glove'  # path of pre-trained word embedding
        elif 'fasttext_wiki' == embedding_name:
            self.word_dim = 300
            self.path_word = 'pretrain/fasttext_wiki'
        elif 'fasttext_crawl' == embedding_name:
            self.word_dim = 300
            self.path_word = 'pretrain/fasttext_crawl'
        elif 'senna' == embedding_name:
            self.word_dim = 50
            self.path_word = 'pretrain/senna'

    def load_embedding(self):
        word2id = dict()  # word to wordID
        word2vec = list()  # wordID to word embedding

        if 'fasttext_wiki' == self.embedding_name:
            if not os.path.exists(os.path.join(self.path_word, FASTTEXT_WIKI_NAME + '_word2id.json')):
                if not os.path.exists(os.path.join(self.path_word, 'wiki-news-300d-1M.vec')):
                    subprocess.call(["bash", "pretrain/download_fasttext_wiki.sh"])
                print("Loading embedding {}...".format(self.embedding_name))
                fasttext_model = KeyedVectors.load_word2vec_format(os.path.join(self.path_word, 'wiki-news-300d-1M.vec'))
                word2id = fasttext_model.key_to_index
                word2vec = np.array([fasttext_model.word_vec(k) for k in word2id.keys()])
                json.dump(word2id, open(os.path.join(self.path_word,FASTTEXT_WIKI_NAME + '_word2id.json'), 'w'))
                np.save(os.path.join(self.path_word,FASTTEXT_WIKI_NAME + '._mat.npy'), word2vec)
            word2id = json.load(open(os.path.join(self.path_word, FASTTEXT_WIKI_NAME + '_word2id.json')))
            word2vec = np.load(os.path.join(self.path_word, FASTTEXT_WIKI_NAME + '._mat.npy'))

        elif 'fasttext_crawl' == self.embedding_name:
            if not os.path.exists(os.path.join(self.path_word, FASTTEXT_CRAWL_NAME + '_word2id.json')):
                if not os.path.exists(os.path.join(self.path_word, 'crawl-300d-2M.vec')):
                    subprocess.call(["bash", "pretrain/download_fasttext_crawl.sh"])
                print("Loading embedding {}...".format(self.embedding_name))
                fasttext_model = KeyedVectors.load_word2vec_format(os.path.join(self.path_word, 'crawl-300d-2M.vec'))
                word2id = fasttext_model.key_to_index
                word2vec = np.array([fasttext_model.word_vec(k) for k in word2id.keys()])
                json.dump(word2id, open(os.path.join(self.path_word,FASTTEXT_CRAWL_NAME + '_word2id.json'), 'w'))
                np.save(os.path.join(self.path_word,FASTTEXT_CRAWL_NAME + '._mat.npy'), word2vec)
            word2id = json.load(open(os.path.join(self.path_word, FASTTEXT_CRAWL_NAME + '_word2id.json')))
            word2vec = np.load(os.path.join(self.path_word, FASTTEXT_CRAWL_NAME + '._mat.npy'))

        elif 'glove' == self.embedding_name:
            if not os.path.exists(os.path.join(self.path_word, GLOVE_NAME + '_mat.npy')):
                subprocess.call(["bash", "pretrain/download_glove.sh"])
            print("Loading embedding {}...".format(self.embedding_name))
            word2id = json.load(open(os.path.join(self.path_word, GLOVE_NAME + '_word2id.json')))
            word2vec = np.load(os.path.join(self.path_word, GLOVE_NAME + '_mat.npy'))

        elif 'senna' == self.embedding_name:
            if not os.path.exists(os.path.join(self.path_word, SENNA_NAME + '_word2id.json')):
                if not os.path.exists(os.path.join(self.path_word, "senna", "embeddings", "embeddings.txt")):
                    subprocess.call(["bash", "pretrain/download_senna.sh"])
                print("Loading embedding {}...".format(self.embedding_name))
                fr = open(os.path.join(self.path_word, 'senna', 'embeddings', 'embeddings.txt'), 'r', encoding='utf-8').readlines()
                w = open(os.path.join(self.path_word, 'senna', 'hash', 'words.lst'), 'r', encoding='utf-8').readlines()
                for i, line in enumerate(fr):
                    line = line.strip().split()
                    if len(line) != self.word_dim:
                        continue
                    word2id[w[i].strip()] = len(word2id)
                    word2vec.append(np.array(line).astype(np.float))

                json.dump(word2id, open(os.path.join(self.path_word,SENNA_NAME + '_word2id.json'), 'w'))
                np.save(os.path.join(self.path_word,SENNA_NAME + '._mat.npy'), word2vec)
            word2id = json.load(open(os.path.join(self.path_word, SENNA_NAME + '_word2id.json')))
            word2vec = np.load(os.path.join(self.path_word, SENNA_NAME + '._mat.npy'))

        return word2id, word2vec

max_stars_repo_name: igorvlnascimento/DeepREF

max_issues_repo_name: igorvlnascimento/DeepREF

max_forks_repo_name: igorvlnascimento/DeepREF


text:
import pytest
from numpy import arange, split, array_equal, empty, newaxis
from bolt import array, ones
from bolt.utils import allclose

def test_chunk(sc):

    x = arange(4*6).reshape(1, 4, 6)
    b = array(x, sc)

    k1, v1 = zip(*b.chunk((2,3))._rdd.sortByKey().collect())
    k2 = ((0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1))
    v2 = [s for m in split(x[0], (2,), axis=0) for s in split(m, (3,), axis=1)]
    assert k1 == k2
    assert all([allclose(m1, m2) for (m1, m2) in zip(v1, v2)])

    k1, v1 = zip(*b.chunk((3,4))._rdd.sortByKey().collect())
    k2 = ((0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1))
    v2 = [s for m in split(x[0], (3,), axis=0) for s in split(m, (4,), axis=1)]
    assert k1 == k2
    assert all([allclose(m1, m2) for (m1, m2) in zip(v1, v2)])

def test_unchunk(sc):

    x = arange(4*6).reshape(1, 4, 6)
    b = array(x, sc)

    assert allclose(b.chunk((2, 3)).unchunk().toarray(), b.toarray())
    assert allclose(b.chunk((3, 4)).unchunk().toarray(), b.toarray())
    assert allclose(b.chunk((4, 6)).unchunk().toarray(), b.toarray())
    assert allclose(b.chunk('0.1').unchunk().toarray(), b.toarray())
    assert allclose(b.chunk().unchunk().toarray(), b.toarray())

    x = arange(4*5*10).reshape(1, 4, 5, 10)
    b = array(x, sc)

    assert allclose(b.chunk((4, 5, 10)).unchunk().toarray(), b.toarray())
    assert allclose(b.chunk((1, 1, 1)).unchunk().toarray(), b.toarray())
    assert allclose(b.chunk((3, 3, 3)).unchunk().toarray(), b.toarray())
    assert allclose(b.chunk((3, 3, 3)).unchunk().toarray(), b.toarray())

    x = arange(4*6).reshape(4, 6)
    b = array(x, sc, (0, 1))

    assert allclose(b.chunk(()).unchunk().toarray(), b.toarray())

    b = array(x, sc, (0,))

    assert allclose(b.chunk((2)).unchunk().toarray(), b.toarray())

def test_keys_to_values(sc):

    x = arange(4*7*9*6).reshape(4, 7, 9, 6)
    b = array(x, sc, (0, 1))
    c = b.chunk((4, 2))

    assert allclose(x, c.keys_to_values((0,)).unchunk().toarray().transpose(1, 0, 2, 3))
    assert allclose(x, c.keys_to_values((1,)).unchunk().toarray())
    assert allclose(x, c.keys_to_values((1,), size=(3,)).unchunk().toarray())
    assert allclose(x, c.keys_to_values((0, 1)).unchunk().toarray())
    assert allclose(x, c.keys_to_values((0, 1), size=(2, 3)).unchunk().toarray())
    assert allclose(x, c.keys_to_values(()).unchunk().toarray())

    b = array(x, sc, range(4))
    c = b.chunk(())

    assert allclose(x, c.keys_to_values((3,)).unchunk().toarray())
    assert allclose(x, c.keys_to_values((0, 1)).unchunk().toarray().transpose(2, 3, 0, 1))

    b = array(x, sc, (0,))
    c = b.chunk((2, 3, 4))

    assert allclose(x, c.keys_to_values((0,)).unchunk().toarray())

def test_values_to_keys(sc):

    x = arange(4*7*9*6).reshape(4, 7, 9, 6)
    b = array(x, sc, (0, 1))
    c = b.chunk((4, 2))

    assert allclose(x, c.values_to_keys((0,)).unchunk().toarray())
    assert allclose(x, c.values_to_keys((1,)).unchunk().toarray().transpose(0, 1, 3, 2))
    assert allclose(x, c.values_to_keys((0, 1)).unchunk().toarray())
    assert allclose(x, c.values_to_keys(()).unchunk().toarray())

    b = array(x, sc, (0,))
    c = b.chunk((2, 3, 4))

    assert allclose(x, c.values_to_keys((0,)).unchunk().toarray())
    assert allclose(x, c.values_to_keys((0, 1)).unchunk().toarray())

def test_padding(sc):

    x = arange(2*2*5*6).reshape(2, 2, 5, 6)
    b = array(x, sc, (0, 1))

    c = b.chunk((2, 2), padding=1)
    chunks = c.tordd().sortByKey().values().collect()
    assert allclose(chunks[0], array([[0, 1, 2], [6, 7, 8], [12, 13, 14]]))
    assert allclose(chunks[1], array([[1, 2, 3, 4], [7, 8, 9, 10], [13, 14, 15, 16]]))
    assert allclose(chunks[4], array([[7, 8, 9, 10], [13, 14, 15, 16], [19, 20, 21, 22], [25, 26, 27, 28]]))
    assert allclose(chunks[6], array([[18, 19, 20], [24, 25, 26]]))

    c = b.chunk((3, 3), padding=(1, 2))
    chunks = c.tordd().sortByKey().values().collect()
    assert allclose(chunks[0], array([[0, 1, 2, 3, 4], [6, 7, 8, 9, 10], [12, 13, 14, 15, 16], [18, 19, 20, 21, 22]]))

    c = b.chunk((2,2), padding=1)
    assert allclose(x, c.unchunk().toarray())
    assert allclose(x, c.keys_to_values((1,)).unchunk().toarray())
    assert allclose(x, c.values_to_keys((0,)).unchunk().toarray())

def test_padding_errors(sc):

        x = arange(2*2*5*6).reshape(2, 2, 5, 6)
        b = array(x, sc, (0, 1))

        with pytest.raises(ValueError):
            c = b.chunk((2, 2), padding=(3, 1))

        with pytest.raises(ValueError):
            c = b.chunk((4, 4), padding=(2, 2))

        with pytest.raises(NotImplementedError):
            c = b.chunk((2, 2), padding=1)
            d = c.map(lambda x: x[:, 0])

def test_map(sc):

    x = arange(4*8*8).reshape(4, 8, 8)
    b = array(x, sc)

    c = b.chunk(size=(4, 8))

    # no change of shape
    def f(x):
        return 2*x

    assert allclose(c.map(f).unchunk().toarray(), f(x))
    assert allclose(c.map(f, value_shape=(4, 8)).unchunk().toarray(), f(x))

    # changing the size of an unchunked axis
    def f(x):
        return x[:, :4]
    def f_local(x):
        return x[:, :, :4]

    assert allclose(c.map(f).unchunk().toarray(), f_local(x))
    assert allclose(c.map(f, value_shape=(4, 4)).unchunk().toarray(), f_local(x))

def test_map_errors(sc):

    x = arange(4*8*8).reshape(4, 8, 8)
    b = array(x, sc)

    c = b.chunk(size=(4, 8))

    # changing the size of a chunked axis
    def f(x):
        return x[:2, :]

    with pytest.raises(ValueError):
        c.map(f)

    with pytest.raises(ValueError):
        c.map(f, value_shape=(2, 8))

    # dropping dimensions
    def f(x):
        return x[0, :]

    with pytest.raises(NotImplementedError):
        c.map(f)

    with pytest.raises(NotImplementedError):
        c.map(f, value_shape=(4,))

def test_map_generic(sc):

    x = arange(2*8*8).reshape(2, 8, 8)
    b = array(x, sc)

    c = b.chunk(size=(8, 5))
    d = c.map_generic(lambda x: [0, 1]).toarray()

    truth = empty(2*1*2, dtype=object)
    for i in range(truth.shape[0]):
        truth[i] = [0, 1]
    truth = truth.reshape(2, 1, 2)

    assert array_equal(d, truth)

def test_properties(sc):

    x = arange(4*6).reshape(1, 4, 6)
    b = array(x, sc)

    assert b.chunk(size=(2, 3)).uniform is True
    assert b.chunk(size=(2, 4)).uniform is False

def test_args(sc):

    x = arange(4*6).reshape(1, 4, 6)
    b = array(x, sc)

    with pytest.raises(ValueError):
        b.chunk(size=(5, 6))

max_stars_repo_name: thunder-project/bolt

max_issues_repo_name: thunder-project/bolt

max_forks_repo_name: thunder-project/bolt