text: from __future__ import absolute_import, division, print_function
  import boost_adaptbx.boost.python as bp
  ext = bp.import_ext("smtbx_refinement_least_squares_ext")
  from smtbx_refinement_least_squares_ext import *
  
  
  import smtbx.refinement.weighting_schemes # import dependency
  from cctbx import xray
  import libtbx.load_env
  from libtbx import adopt_optional_init_args
  from scitbx import linalg
  from scitbx.lstbx import normal_eqns
  from scitbx.array_family import flex
  from smtbx.structure_factors import direct
  from smtbx.refinement.restraints import origin_fixing_restraints
  import math
  
  def crystallographic_ls_class(non_linear_ls_with_separable_scale_factor=None):
    """ Construct a class for crystallographic L.S. based on the given engine
    """
    def get_base_class(non_linear_ls_with_separable_scale_factor):
      base_class = non_linear_ls_with_separable_scale_factor
      if not base_class:
        try:
          from fast_linalg import env
          if env.initialised:
            base_class = normal_eqns.non_linear_ls_with_separable_scale_factor_BLAS_3
          else:
            base_class = normal_eqns.non_linear_ls_with_separable_scale_factor_BLAS_2
        except Exception:
          base_class = normal_eqns.non_linear_ls_with_separable_scale_factor_BLAS_2
      #print("Chosen: " + str(base_class))
      return base_class
  
    class klass(get_base_class(non_linear_ls_with_separable_scale_factor)):
  
      non_linear_ls_engine = get_base_class(non_linear_ls_with_separable_scale_factor)
  
      default_weighting_scheme = mainstream_shelx_weighting
      weighting_scheme = "default"
      origin_fixing_restraints_type = (
        origin_fixing_restraints.atomic_number_weighting)
      f_mask = None
      restraints_manager=None
      n_restraints = None
      initial_scale_factor = None
      may_parallelise = False
      use_openmp = False
  
      def __init__(self, observations, reparametrisation,
                   one_h_linearisation=None, **kwds):
        super(klass, self).__init__(reparametrisation.n_independents)
        self.observations = observations
        self.reparametrisation = reparametrisation
        adopt_optional_init_args(self, kwds)
        if self.f_mask is not None:
          assert self.f_mask.size() == observations.fo_sq.size()
        self.one_h_linearisation = one_h_linearisation
        if not self.one_h_linearisation:
          self.one_h_linearisation = direct.f_calc_modulus_squared(
            self.xray_structure)
        if self.weighting_scheme == "default":
          self.weighting_scheme = self.default_weighting_scheme()
        self.origin_fixing_restraint = self.origin_fixing_restraints_type(
          self.xray_structure.space_group())
        self.taken_step = None
        self.restraints_normalisation_factor = None
  
      @property
      def xray_structure(self):
        return self.reparametrisation.structure
  
      @property
      def twin_fractions(self):
        return self.reparametrisation.twin_fractions
  
      def build_up(self, objective_only=False):
        if self.f_mask is not None:
          f_mask = self.f_mask.data()
        else:
          f_mask = flex.complex_double()
  
        extinction_correction = self.reparametrisation.extinction
        if extinction_correction is None:
          extinction_correction = xray.dummy_extinction_correction()
  
        def build_normal_eqns(scale_factor, weighting_scheme, objective_only):
          return ext.build_normal_equations(
            self,
            self.observations,
            f_mask,
            weighting_scheme,
            scale_factor,
            self.one_h_linearisation,
            self.reparametrisation.jacobian_transpose_matching_grad_fc(),
            extinction_correction,
            objective_only,
            self.may_parallelise,
            self.use_openmp)
  
        if not self.finalised: #i.e. never been called
          self.reparametrisation.linearise()
          self.reparametrisation.store()
          scale_factor = self.initial_scale_factor
          if scale_factor is None: # we haven't got one from previous refinement
            result = build_normal_eqns(scale_factor=None,
                                       weighting_scheme=sigma_weighting(),
                                       objective_only=True)
            scale_factor = self.scale_factor()
        else: # use scale factor from the previous step
          scale_factor = self.scale_factor()
  
        self.reset()
        result = build_normal_eqns(scale_factor,
                                   self.weighting_scheme,
                                   objective_only)
        self.f_calc = self.observations.fo_sq.array(
          data=result.f_calc(), sigmas=None)
        self.fc_sq = self.observations.fo_sq.array(
          data=result.observables(), sigmas=None)\
            .set_observation_type_xray_intensity()
        self.weights = result.weights()
        self.objective_data_only = self.objective()
        self.chi_sq_data_only = self.chi_sq()
        if self.restraints_manager is not None:
          # Here we determine a normalisation factor to place the restraints on the
          # same scale as the average residual. This is the normalisation
          # factor suggested in Giacovazzo and similar to that used by shelxl.
          # (shelx manual, page 5-1).
          # The factor 2 comes from the fact that we minimize 1/2 sum w delta^2
          if self.restraints_normalisation_factor is None:
            self.restraints_normalisation_factor \
                = 2 * self.objective_data_only/(self.n_equations-self.n_parameters)
          linearised_eqns = self.restraints_manager.build_linearised_eqns(
            self.xray_structure, self.reparametrisation.parameter_map())
          jacobian = \
            self.reparametrisation.jacobian_transpose_matching(
              self.reparametrisation.mapping_to_grad_fc_all).transpose()
          self.reduced_problem().add_equations(
            linearised_eqns.deltas,
            linearised_eqns.design_matrix * jacobian,
            linearised_eqns.weights * self.restraints_normalisation_factor,
            optimise_for_tall_matrix=False)
          self.n_restraints = linearised_eqns.n_restraints()
          self.chi_sq_data_and_restraints = self.chi_sq()
        if not objective_only:
          self.origin_fixing_restraint.add_to(
            self.step_equations(),
            self.reparametrisation.jacobian_transpose_matching_grad_fc(),
            self.reparametrisation.asu_scatterer_parameters)
  
      def parameter_vector_norm(self):
        return self.reparametrisation.norm_of_independent_parameter_vector
  
      def scale_factor(self): return self.optimal_scale_factor()
  
      def step_forward(self):
        self.reparametrisation.apply_shifts(self.step())
        self.reparametrisation.linearise()
        self.reparametrisation.store()
        self.taken_step = self.step().deep_copy()
  
      def step_backward(self):
        self.reparametrisation.apply_shifts(-self.taken_step)
        self.reparametrisation.linearise()
        self.reparametrisation.store()
        self.taken_step = None
  
      def goof(self):
        return math.sqrt(self.chi_sq_data_only)
  
      def restrained_goof(self):
        if self.restraints_manager is None:
          return self.goof()
        return math.sqrt(self.chi_sq_data_and_restraints)
  
      def wR2(self, cutoff_factor=None):
        if cutoff_factor is None:
          return math.sqrt(2*self.objective_data_only)
        fo_sq = self.observations.fo_sq
        strong = fo_sq.data() >= cutoff_factor*fo_sq.sigmas()
        fo_sq = fo_sq.select(strong)
        fc_sq = self.fc_sq.select(strong)
        wght = self.weights.select(strong)
        fc_sq = fc_sq.data()
        fo_sq = fo_sq.data()
        fc_sq *= self.scale_factor()
        wR2 = flex.sum(wght*flex.pow2((fo_sq-fc_sq)))/flex.sum(wght*flex.pow2(fo_sq))
        return math.sqrt(wR2)
  
      def r1_factor(self, cutoff_factor=None):
        fo_sq = self.observations.fo_sq
        if cutoff_factor is not None:
          strong = fo_sq.data() >= cutoff_factor*fo_sq.sigmas()
          fo_sq = fo_sq.select(strong)
          fc_sq = self.fc_sq.select(strong)
        else:
          fc_sq = self.fc_sq
        f_obs = fo_sq.f_sq_as_f()
        f_calc = fc_sq.f_sq_as_f()
        R1 = f_obs.r1_factor(f_calc,
          scale_factor=math.sqrt(self.scale_factor()), assume_index_matching=True)
        return R1, f_obs.size()
  
      def covariance_matrix(self,
                            jacobian_transpose=None,
                            normalised_by_goof=True):
        """ The columns of the jacobian_transpose determine which crystallographic
            parameters appear in the covariance matrix.
            If jacobian_transpose is None, then the covariance matrix returned will
            be that for the independent L.S. parameters.
        """
        if not self.step_equations().solved:
          self.solve()
        cov = linalg.inverse_of_u_transpose_u(
          self.step_equations().cholesky_factor_packed_u())
        cov /= self.sum_w_yo_sq()
        if jacobian_transpose is not None:
          cov = jacobian_transpose.self_transpose_times_symmetric_times_self(cov)
        if normalised_by_goof: cov *= self.restrained_goof()**2
        return cov
  
      def covariance_matrix_and_annotations(self):
        jac_tr = self.reparametrisation.jacobian_transpose_matching_grad_fc()
        return covariance_matrix_and_annotations(
          self.covariance_matrix(jacobian_transpose=jac_tr),
          self.reparametrisation.component_annotations)
  
    return klass
  
  
  def crystallographic_ls(
    observations, reparametrisation,
    non_linear_ls_with_separable_scale_factor=None,
    may_parallelise=True,
    **kwds):
    return crystallographic_ls_class(non_linear_ls_with_separable_scale_factor
                                     )(observations, reparametrisation,
                                       may_parallelise=may_parallelise, **kwds)
  
  
  class covariance_matrix_and_annotations(object):
  
    def __init__(self, covariance_matrix, annotations):
      """ The covariance matrix is assumed to be a symmetric matrix stored as a
          packed upper diagonal matrix.
      """
      self.matrix = covariance_matrix
      self.annotations = annotations
      self._2_n_minus_1 = 2*len(self.annotations)-1 # precompute for efficiency
  
    def __call__(self, i, j):
      return self.matrix[i*(self._2_n_minus_1-i)//2 + j]
  
    def variance_of(self, annotation):
      i = self.annotations.index(annotation)
      return self(i, i)
  
    def covariance_of(self, annotation_1, annotation_2):
      i = self.annotations.index(annotation_1)
      j = self.annotations.index(annotation_2)
      if j > i:
        i, j = j, i
      return self(i, j)
  
    def diagonal(self):
      return self.matrix.matrix_packed_u_diagonal()
  