text: #!/usr/bin/env python3
  # -*- coding: utf-8 -*-
  """Module holding the classes for IRGN Optimization without streaming."""
  from __future__ import division
  import time
  import numpy as np
  
  from pkg_resources import resource_filename
  import pyopencl.array as clarray
  import h5py
  
  import pyqmri.operator as operator
  import pyqmri.solver as optimizer
  from pyqmri._helper_fun import CLProgram as Program
  from pyqmri._helper_fun import _utils as utils
  
  
  class IRGNOptimizer:
      """Main IRGN Optimization class.
  
      This Class performs IRGN Optimization either with TGV or TV regularization.
  
      Parameters
      ----------
        par : dict
          A python dict containing the necessary information to
          setup the object. Needs to contain the number of slices (NSlice),
          number of scans (NScan), image dimensions (dimX, dimY), number of
          coils (NC), sampling points (N) and read outs (NProj)
          a PyOpenCL queue (queue) and the complex coil
          sensitivities (C).
        model : pyqmri.model
          Which model should be used for fitting.
          Expects a pyqmri.model instance.
        trafo : int, 1
          Select radial (1, default) or cartesian (0) sampling for the fft.
        imagespace : bool, false
          Perform the fitting from k-space data (false, default) or from the
          image series (true)
        SMS : int, 0
          Select if simultaneous multi-slice acquisition was used (1) or
          standard slice-by-slice acquisition was done (0, default).
        reg_type : str, "TGV"
          Select between "TGV" (default) or "TV" regularization.
        config : str, ''
          Name of config file. If empty, default config file will be generated.
        streamed : bool, false
          Select between standard reconstruction (false)
          or streamed reconstruction (true) for large volumetric data which
          does not fit on the GPU memory at once.
        DTYPE : numpy.dtype, numpy.complex64
          Complex working precission.
        DTYPE_real : numpy.dtype, numpy.float32
          Real working precission.
  
      Attributes
      ----------
        par : dict
          A python dict containing the necessary information to
          setup the object. Needs to contain the number of slices (NSlice),
          number of scans (NScan), image dimensions (dimX, dimY), number of
          coils (NC), sampling points (N) and read outs (NProj)
          a PyOpenCL queue (queue) and the complex coil
          sensitivities (C).
        gn_res : list of floats
          The residual values for of each Gauss-Newton step. Each iteration
          appends its value to the list.
        irgn_par : dict
          The parameters read from the config file to guide the IRGN
          optimization process
      """
  
      def __init__(self, par, model, trafo=1, imagespace=False, SMS=0,
                   reg_type='TGV', config='', streamed=False,
                   DTYPE=np.complex64, DTYPE_real=np.float32):
          self.par = par
          self.gn_res = []
          self.irgn_par = utils.read_config(config, reg_type)
          utils.save_config(self.irgn_par, par["outdir"], reg_type)
          num_dev = len(par["num_dev"])
          self._fval_old = 0
          self._fval = 0
          self._fval_init = 0
          self._ctx = par["ctx"]
          self._queue = par["queue"]
          self._model = model
          self._reg_type = reg_type
          self._prg = []
  
          self._DTYPE = DTYPE
          self._DTYPE_real = DTYPE_real
  
          self._streamed = streamed
          self._imagespace = imagespace
          self._SMS = SMS
          if streamed and par["NSlice"]/(num_dev*par["par_slices"]) < 2:
              raise ValueError(
                  "Number of Slices devided by parallel "
                  "computed slices and devices needs to be larger two.\n"
                  "Current values are %i total Slices, %i parallel slices and "
                  "%i compute devices."
                  % (par["NSlice"], par["par_slices"], num_dev))
          if streamed and par["NSlice"] % par["par_slices"]:
              raise ValueError(
                  "Number of Slices devided by parallel "
                  "computed slices needs to be an integer.\n"
                  "Current values are %i total Slices with %i parallel slices."
                  % (par["NSlice"], par["par_slices"]))
          if DTYPE == np.complex128:
              if streamed:
                  kernname = 'kernels/OpenCL_Kernels_double_streamed.c'
              else:
                  kernname = 'kernels/OpenCL_Kernels_double.c'
              for j in range(num_dev):
                  self._prg.append(Program(
                      self._ctx[j],
                      open(
                          resource_filename(
                              'pyqmri', kernname)
                          ).read()))
          else:
              if streamed:
                  kernname = 'kernels/OpenCL_Kernels_streamed.c'
              else:
                  kernname = 'kernels/OpenCL_Kernels.c'
              for j in range(num_dev):
                  self._prg.append(Program(
                      self._ctx[j],
                      open(
                          resource_filename(
                              'pyqmri', kernname)).read()))
  
          if imagespace:
              self._coils = []
              self.sliceaxis = 1
              if self._streamed:
                  self._data_trans_axes = (1, 0, 2, 3)
                  self._grad_trans_axes = (2, 0, 1, 3, 4)
          else:
              if SMS:
                  self._data_shape = (par["NScan"], par["NC"],
                                      par["packs"]*par["numofpacks"],
                                      par["Nproj"], par["N"])
              else:
                  if par["is3D"] and trafo:
                      self._data_shape = (par["NScan"], par["NC"],
                                          1, par["Nproj"], par["N"])
                  else:
                      self._data_shape = (par["NScan"], par["NC"],
                                          par["NSlice"], par["Nproj"], par["N"])
              if self._streamed:
                  self._data_trans_axes = (2, 0, 1, 3, 4)
                  self._grad_trans_axes = (2, 0, 1, 3, 4)
                  self._coils = np.require(
                      np.swapaxes(par["C"], 0, 1), requirements='C',
                      dtype=DTYPE)
  
                  if SMS:
                      self._data_shape = (par["packs"]*par["numofpacks"],
                                          par["NScan"],
                                          par["NC"], par["dimY"], par["dimX"])
                      self._data_shape_T = (par["NScan"], par["NC"],
                                            par["packs"]*par["numofpacks"],
                                            par["dimY"], par["dimX"])
                      self._expdim_dat = 1
                      self._expdim_C = 0
                  else:
                      self._data_shape = (par["NSlice"], par["NScan"], par["NC"],
                                          par["Nproj"], par["N"])
                      self._data_shape_T = self._data_shape
                      self._expdim_dat = 2
                      self._expdim_C = 1
              else:
                  self._coils = clarray.to_device(self._queue[0],
                                                  self.par["C"])
  
          self._MRI_operator, self._FT = operator.Operator.MRIOperatorFactory(
              par,
              self._prg,
              DTYPE,
              DTYPE_real,
              trafo,
              imagespace,
              SMS,
              streamed
              )
  
          self._grad_op, self._symgrad_op, self._v = self._setupLinearOps(
              DTYPE,
              DTYPE_real)
  
          if not reg_type == "H1":
              self._pdop = optimizer.PDBaseSolver.factory(
                  self._prg,
                  self._queue,
                  self.par,
                  self.irgn_par,
                  self._fval_init,
                  self._coils,
                  linops=(self._MRI_operator, self._grad_op, self._symgrad_op),
                  model=model,
                  reg_type=self._reg_type,
                  SMS=self._SMS,
                  streamed=self._streamed,
                  imagespace=self._imagespace,
                  DTYPE=DTYPE,
                  DTYPE_real=DTYPE_real
                  )
          else:
              self._pdop = optimizer.CGSolver_H1(
                  self._prg,
                  self._queue,
                  self.par,
                  self.irgn_par,
                  self._coils,
                  linops=(self._MRI_operator, self._grad_op)
                  )
  
          self._gamma = None
          self._delta = None
          self._omega = None
          self._step_val = None
          self._modelgrad = None
  
      def _setupLinearOps(self, DTYPE, DTYPE_real):
          grad_op = operator.Operator.GradientOperatorFactory(
              self.par,
              self._prg,
              DTYPE,
              DTYPE_real,
              self._streamed)
          symgrad_op = None
          v = None
          if self._reg_type == 'TGV':
              symgrad_op = operator.Operator.SymGradientOperatorFactory(
                  self.par,
                  self._prg,
                  DTYPE,
                  DTYPE_real,
                  self._streamed)
              v = np.zeros(
                  ([self.par["unknowns_TGV"], self.par["NSlice"],
                    self.par["dimY"], self.par["dimX"], 4]),
                  dtype=DTYPE)
              if self._streamed:
                  v = np.require(np.swapaxes(v, 0, 1), requirements='C')
          return grad_op, symgrad_op, v
  
      def execute(self, data):
          """Start the IRGN optimization.
  
          This method performs iterative regularized Gauss-Newton optimization
          and calls the inner loop after precomputing the current linearization
          point. Results of the fitting process are saved after each
          linearization step to the output folder.
  
          Parameters
          ----------
            data : numpy.array
              the data to perform optimization/fitting on.
          """
          # self.irgn_par["lambd"] *= (
          #                             (self.par["SNR_est"]))
          self._gamma = self.irgn_par["gamma"]
          self._delta = self.irgn_par["delta"]
          self._omega = self.irgn_par["omega"]
  
          iters = self.irgn_par["start_iters"]
          result = np.copy(self._model.guess)
  
          if "inc" in self.irgn_par.keys():
              inc = int(self.irgn_par["inc"])
          else:
              inc = 2
  
          if self._streamed:
              data = np.require(
                  np.transpose(data, self._data_trans_axes),
                  requirements='C')
  
          self._step_val = np.nan_to_num(self._model.execute_forward(result))
  
          if self._streamed:
              if self._SMS is False:
                  self._step_val = np.require(
                      np.swapaxes(self._step_val, 0, 1), requirements='C')
  
          for ign in range(self.irgn_par["max_gn_it"]):
              start = time.time()
              self._modelgrad = np.nan_to_num(
                  self._model.execute_gradient(result))
  
              self._balanceModelGradients(result, ign)
              self._updateIRGNRegPar(ign)
              # self._pdop._grad_op.updateRatio(
              #     self.irgn_par["gamma"])
              # if self._reg_type == 'TGV':
              #     self._pdop._symgrad_op.updateRatio(
              #         2*self.irgn_par["gamma"])
  
              self._step_val = np.nan_to_num(self._model.execute_forward(result))
  
              if self._streamed:
                  if self._SMS is False:
                      self._step_val = np.require(
                          np.swapaxes(self._step_val, 0, 1), requirements='C')
                  self._modelgrad = np.require(
                      np.transpose(self._modelgrad, self._grad_trans_axes),
                      requirements='C')
                  self._pdop.model = self._model
                  self._pdop.modelgrad = self._modelgrad
                  self._pdop.jacobi = np.sum(
                      np.abs(self._modelgrad)**2, 2).astype(self._DTYPE_real)
                  self._pdop.jacobi[self._pdop.jacobi == 0] = 1e-8
              else:
                  _jacobi = np.sum(
                      np.abs(
                          self._modelgrad)**2, 1).astype(self._DTYPE_real)
                  _jacobi[_jacobi == 0] = 1e-8
                  self._modelgrad = clarray.to_device(
                      self._queue[0],
                      self._modelgrad)
                  self._pdop.model = self._model
                  self._pdop.modelgrad = self._modelgrad
                  self._pdop.jacobi = clarray.to_device(
                      self._queue[0],
                      _jacobi)
  
              self._pdop.updateRegPar(self.irgn_par)
  
              result = self._irgnSolve3D(result, iters, data, ign)
  
              iters = np.fmin(iters * inc, self.irgn_par["max_iters"])
  
              end = time.time() - start
              self.gn_res.append(self._fval)
              print("-" * 75)
              print("GN-Iter: %d  Elapsed time: %f seconds" % (ign, end))
              print("-" * 75)
              self._fval_old = self._fval
              self._saveToFile(ign, self._model.rescale(result)["data"])
          self._calcResidual(result, data, ign+1)
  
      def _updateIRGNRegPar(self, ign):
          try:
              self.irgn_par["delta"] = np.minimum(
                  self._delta
                  * self.irgn_par["delta_inc"]**ign,
                  self.irgn_par["delta_max"])*self.irgn_par["lambd"]
          except OverflowError:
              self.irgn_par["delta"] = np.minimum(
                  np.finfo(self._delta).max,
                  self.irgn_par["delta_max"])
          self.irgn_par["gamma"] = np.maximum(
              self._gamma * self.irgn_par["gamma_dec"]**ign,
              self.irgn_par["gamma_min"])/self.irgn_par["lambd"]
          self.irgn_par["omega"] = np.maximum(
              self._omega * self.irgn_par["omega_dec"]**ign,
              self.irgn_par["omega_min"])/self.irgn_par["lambd"]
  
      def _balanceModelGradients(self, result, ign):
          scale = np.reshape(
              self._modelgrad,
              (self.par["unknowns"],
               self.par["NScan"] * self.par["NSlice"] *
               self.par["dimY"] * self.par["dimX"]))
          scale = np.linalg.norm(scale, axis=-1)
          print("Initial Norm: ", np.linalg.norm(scale))
          print("Initial Ratio: ", scale)
          if ign == 0:
              scale /= 100/np.sqrt(self.par["unknowns"])
          else:
              scale /= np.linalg.norm(scale)/np.sqrt(self.par["unknowns"])
          scale = 1 / scale
          scale[~np.isfinite(scale)] = 1
          for uk in range(self.par["unknowns"]):
              self._model.constraints[uk].update(scale[uk])
              result[uk, ...] *= self._model.uk_scale[uk]
              self._modelgrad[uk] /= self._model.uk_scale[uk]
              self._model.uk_scale[uk] *= scale[uk]
              result[uk, ...] /= self._model.uk_scale[uk]
              self._modelgrad[uk] *= self._model.uk_scale[uk]
          scale = np.reshape(
              self._modelgrad,
              (self.par["unknowns"],
               self.par["NScan"] * self.par["NSlice"] *
               self.par["dimY"] * self.par["dimX"]))
          scale = np.linalg.norm(scale, axis=-1)
          print("Norm after rescale: ", np.linalg.norm(scale))
          print("Ratio after rescale: ", scale)
  
  ###############################################################################
  # New .hdf5 save files ########################################################
  ###############################################################################
      def _saveToFile(self, myit, result):
          f = h5py.File(self.par["outdir"]+"output_" + self.par["fname"] + ".h5",
                        "a")
          if self._reg_type == 'TGV':
              f.create_dataset("tgv_result_iter_"+str(myit), result.shape,
                               dtype=self._DTYPE, data=result)
              f.attrs['res_tgv_iter_'+str(myit)] = self._fval
          else:
              f.create_dataset("tv_result_"+str(myit), result.shape,
                               dtype=self._DTYPE, data=result)
              f.attrs['res_tv_iter_'+str(myit)] = self._fval
          f.close()
  
  ###############################################################################
  # Precompute constant terms of the GN linearization step ######################
  # input: linearization point x ################################################
  # iters: number of innner iterations iters ####################################
  # data: the input data ########################################################
  # TV: bool to switch between TV (1) and TGV (0) regularization ################
  # output: optimal value of x for the inner GN step ############################
  ###############################################################################
  ###############################################################################
      def _irgnSolve3D(self, x, iters, data, GN_it):
          b = self._calcResidual(x, data, GN_it)
          if self._streamed:
              x = np.require(np.swapaxes(x, 0, 1), requirements='C')
              res = data - b + self._MRI_operator.fwdoop(
                  [[x, self._coils, self._modelgrad]])
          else:
              tmpx = clarray.to_device(self._queue[0], x)
              res = data - b + self._MRI_operator.fwdoop(
                  [tmpx, self._coils, self._modelgrad]).get()
              del tmpx
  
          tmpres = self._pdop.run((x, self._v), res, iters)
          for key in tmpres:
              if key == 'x':
                  if isinstance(tmpres[key], np.ndarray):
                      x = tmpres["x"]
                  else:
                      x = tmpres["x"].get()
              if key == 'v':
                  if isinstance(tmpres[key], np.ndarray):
                      self._v = tmpres["v"]
                  else:
                      self._v = tmpres["v"].get()
          if self._streamed:
              x = np.require(np.swapaxes(x, 0, 1), requirements='C')
  
          return x
  
      def _calcResidual(self, x, data, GN_it):
          if self._streamed:
              b, grad, sym_grad = self._calcFwdGNPartStreamed(x)
              grad_tv = grad[:, :self.par["unknowns_TGV"]]
              grad_H1 = grad[:, self.par["unknowns_TGV"]:]
          else:
              b, grad, sym_grad = self._calcFwdGNPartLinear(x)
              grad_tv = grad[:self.par["unknowns_TGV"]]
              grad_H1 = grad[self.par["unknowns_TGV"]:]
          del grad
  
          datacost = 1 / 2 * np.linalg.norm(data - b)**2
          # L2Cost = np.linalg.norm(x)/(2.0*self.irgn_par["delta"])
          if self._reg_type == 'TV':
              regcost = self.irgn_par["gamma"] * \
                  np.sum(np.abs(grad_tv))
          elif self._reg_type == 'TGV':
              regcost = self.irgn_par["gamma"] * np.sum(
                    np.abs(grad_tv -
                           self._v)) + self.irgn_par["gamma"] * 2 * np.sum(
                               np.abs(sym_grad))
              del sym_grad
          else:
              regcost = self.irgn_par["gamma"]/2 * \
                  np.linalg.norm(grad_tv)**2
  
          self._fval = (datacost +
                        regcost +
                         # L2Cost +
                        self.irgn_par["omega"] / 2 *
                        np.linalg.norm(grad_H1.flatten())**2)
          del grad_tv, grad_H1
  
          if GN_it == 0:
              self._fval_init = self._fval
              self._pdop.setFvalInit(self._fval)
  
          print("-" * 75)
          # print("Initial Cost: %f" % (self._fval_init))
          print("Costs of Data: %f" % (1e3*datacost / self._fval_init))
          print("Costs of T(G)V: %f" % (1e3*regcost / self._fval_init))
          # print("Costs of L2 Term: %f" % (1e3*L2Cost / self._fval_init))
          print("-" * 75)
          print("Function value at GN-Step %i: %f" %
                (GN_it, 1e3*self._fval / self._fval_init))
          print("-" * 75)
          return b
  
      def _calcFwdGNPartLinear(self, x):
          if self._imagespace is False:
              b = clarray.zeros(self._queue[0],
                                self._data_shape,
                                dtype=self._DTYPE)
              self._FT.FFT(b, clarray.to_device(
                  self._queue[0],
                  (self._step_val[:, None, ...] *
                   self.par["C"]))).wait()
              b = b.get()
          else:
              b = self._step_val
  
          x = clarray.to_device(self._queue[0], np.require(x, requirements="C"))
          grad = clarray.to_device(self._queue[0],
                                   np.zeros(x.shape+(4,), dtype=self._DTYPE))
          self._grad_op.fwd(
              grad,
              x,
              wait_for=grad.events +
              x.events).wait()
          x = x.get()
          grad = grad.get()
          sym_grad = None
          if self._reg_type == 'TGV':
              v = clarray.to_device(self._queue[0], self._v)
              sym_grad = clarray.to_device(self._queue[0],
                                           np.zeros(x.shape+(8,),
                                                    dtype=self._DTYPE))
  
              self._symgrad_op.fwd(
                  sym_grad,
                  v,
                  wait_for=sym_grad.events +
                  v.events).wait()
              sym_grad = sym_grad.get()
  
          return b, grad, sym_grad
  
      def _calcFwdGNPartStreamed(self, x):
          x = np.require(np.swapaxes(x, 0, 1), requirements='C')
          if self._imagespace is False:
              b = np.zeros(self._data_shape_T, dtype=self._DTYPE)
              if self._SMS is True:
                  self._MRI_operator.FTstr.eval(
                      [b],
                      [[np.expand_dims(self._step_val, self._expdim_dat) *
                        np.expand_dims(self.par["C"], self._expdim_C)]])
                  b = np.require(
                      np.transpose(
                          b,
                          self._data_trans_axes),
                      requirements='C')
              else:
                  self._MRI_operator.FTstr.eval(
                      [b],
                      [[np.expand_dims(self._step_val, self._expdim_dat) *
                        np.expand_dims(self._coils, self._expdim_C)]])
          else:
              b = self._step_val
          grad = np.zeros(x.shape+(4,), dtype=self._DTYPE)
          self._grad_op.fwd([grad], [[x]])
  
          sym_grad = None
          if self._reg_type == 'TGV':
              sym_grad = np.zeros(x.shape+(8,), dtype=self._DTYPE)
              self._symgrad_op.fwd([sym_grad], [[self._v]])
  
          return b, grad, sym_grad
  