<commit_before>
#using CelesteTypes.getids
#using CelesteTypes.HessianEntry

export multiply_sf!, add_scaled_sfs!, combine_sfs!, add_sources_sf!

@doc """
A function value and its derivative with respect to its arguments.

Attributes:
  v:  The value
  d:  The derivative with respect to each variable in
      P-dimensional VariationalParams for each of S celestial objects
      in a local_P x local_S matrix.
  h:  The second derivative with respect to each variational parameter,
      in the same format as d.  This is used for the full Hessian
      with respect to all the sources.
  hs: An array of per-source Hessians.  This will generally be reserved
      for the Hessian of brightness values that depend only on one source.
""" ->
type SensitiveFloat{ParamType <: CelesteTypes.ParamSet, NumType <: Number}
    v::NumType
    d::Matrix{NumType} # local_P x local_S
    # h is ordered so that p changes fastest.  For example, the indices
    # of a column of h correspond to the indices of d's stacked columns.
    h::Matrix{NumType} # (local_P * local_S) x (local_P * local_S)
    ids::ParamType
end


#########################################################

@doc """
Set a SensitiveFloat's hessian term, maintaining symmetry.
""" ->
function set_hess!{ParamType <: CelesteTypes.ParamSet, NumType <: Number}(
    sf::SensitiveFloat{ParamType, NumType},
    i::Int64, j::Int64, v::NumType)
  i != j ?
    sf.h[i, j] = sf.h[j, i] = v:
    sf.h[i, j] = v
end

function zero_sensitive_float{ParamType <: CelesteTypes.ParamSet}(
  ::Type{ParamType}, NumType::DataType, local_S::Int64)
    local_P = length(ParamType)
    d = zeros(NumType, local_P, local_S)
    h = zeros(NumType, local_P * local_S, local_P * local_S)
    SensitiveFloat{ParamType, NumType}(
      zero(NumType), d, h, getids(ParamType))
end

function zero_sensitive_float{ParamType <: CelesteTypes.ParamSet}(
  ::Type{ParamType}, NumType::DataType)
    # Default to a single source.
    zero_sensitive_float(ParamType, NumType, 1)
end

function clear!{ParamType <: CelesteTypes.ParamSet, NumType <: Number}(
  sp::SensitiveFloat{ParamType, NumType})
    sp.v = zero(NumType)
    fill!(sp.d, zero(NumType))
    fill!(sp.h, zero(NumType))
end

# If no type is specified, default to using Float64.
function zero_sensitive_float{ParamType <: CelesteTypes.ParamSet}(
  param_arg::Type{ParamType}, local_S::Int64)
    zero_sensitive_float(param_arg, Float64, local_S)
end

function zero_sensitive_float{ParamType <: CelesteTypes.ParamSet}(
  param_arg::Type{ParamType})
    zero_sensitive_float(param_arg, Float64, 1)
end

function +(sf1::SensitiveFloat, sf2::SensitiveFloat)
  S = size(sf1.d)[2]

  # Simply asserting equality of the ids doesn't work for some reason.
  @assert typeof(sf1.ids) == typeof(sf2.ids)
  @assert length(sf1.ids) == length(sf2.ids)
  [ @assert size(sf1.h[s]) == size(sf2.h[s]) for s=1:S ]

  @assert size(sf1.d) == size(sf2.d)

  sf3 = deepcopy(sf1)
  sf3.v = sf1.v + sf2.v
  sf3.d = sf1.d + sf2.d
  sf3.h = sf1.h + sf2.h

  sf3
end


@doc """
Updates sf1 in place with sf1 * sf2.

ids1 and ids2 are the ids for which sf1 and sf2 have nonzero derivatives,
respectively.
""" ->
function multiply_sf!{ParamType <: CelesteTypes.ParamSet, NumType <: Number}(
    sf1::SensitiveFloat{ParamType, NumType},
    sf2::SensitiveFloat{ParamType, NumType};
    ids1::Array{Int64}=collect(1:length(sf1.ids)),
    ids2::Array{Int64}=collect(1:length(sf2.ids)),
    calculate_hessian::Bool=true)

  S = size(sf1.d)[2]
  @assert S == 1 # For now this is only for the brightness calculations.
  # TODO: replace this with combine_sfs!.

  # You have to do this in the right order to not overwrite needed terms.

  if calculate_hessian
    # Second derivate terms involving second derivates.
    sf1.h[:, :] = sf1.v * sf2.h + sf2.v * sf1.h
    sf1.h[:, :] += sf1.d[:] * sf2.d[:]' + sf2.d[:] * sf1.d[:]'
  end
  sf1.d[:, :] = sf2.v * sf1.d + sf1.v * sf2.d

  sf1.v = sf1.v * sf2.v
end


@doc """
Factor out the hessian part of combine_sfs! to help the compiler.

TODO: I think this is a red herring and this can be put back in
""" ->
function combine_sfs_hessian!{ParamType <: CelesteTypes.ParamSet, NumType <: Number}(
    sf1::SensitiveFloat{ParamType, NumType},
    sf2::SensitiveFloat{ParamType, NumType},
    sf_result::SensitiveFloat{ParamType, NumType},
    g_d::Vector{NumType}, g_h::Matrix{NumType})

  const use_blas = false

  if use_blas
    # Chain rule for second derivatives.
    # BLAS for
    # sf_result.h[:, :] = g_d[1] * sf1.h + g_d[2] * sf2.h
    BLAS.blascopy!(prod(size(sf_result.h)), sf1.h, 1, sf_result.h, 1);
    BLAS.scal!(prod(size(sf_result.h)), g_d[1], sf_result.h, 1);
    BLAS.axpy!(g_d[2], sf2.h, sf_result.h)

    # BLAS for
    # sf_result.h[:, :] +=
    #   g_h[1, 1] * sf1.d[:] * sf1.d[:]' +
    #   g_h[2, 2] * sf2.d[:] * sf2.d[:]' +
    #   g_h[1, 2] * (sf1.d[:] * sf2.d[:]' + sf2.d[:] * sf1.d[:]')
    sf1d = sf1.d[:];
    sf2d = sf2.d[:];
    BLAS.ger!(g_h[1, 1], sf1d, sf1d, sf_result.h);
    BLAS.ger!(g_h[2, 2], sf2d, sf2d, sf_result.h);
    BLAS.ger!(g_h[1, 2], sf1d, sf2d, sf_result.h);
    BLAS.ger!(g_h[1, 2], sf2d, sf1d, sf_result.h);
  else
    p1, p2 = size(sf_result.h)
    for ind1 = 1:p1, ind2 = 1:p2
      sf_result.h[ind1, ind2] =
        g_d[1] * sf1.h[ind1, ind2] + g_d[2] * sf2.h[ind1, ind2] +
        g_h[1, 1] * sf1.d[ind1] * sf1.d[ind2] +
        g_h[2, 2] * sf2.d[ind1] * sf2.d[ind2] +
        g_h[1, 2] * (sf1.d[ind1] * sf2.d[ind2] + sf2.d[ind1] * sf1.d[ind2])
    end
  end
end


@doc """
Updates sf_result in place with g(sf1, sf2), where
g_d = (g_1, g_2) is the gradient of g and
g_h = (g_11, g_12; g_12, g_22) is the hessian of g,
each evaluated at (sf1, sf2).

The result is stored in sf_result.  The order is done in such a way that
it can overwrite sf1 or sf2 and still be accurate.
""" ->
function combine_sfs!{ParamType <: CelesteTypes.ParamSet, NumType <: Number}(
    sf1::SensitiveFloat{ParamType, NumType},
    sf2::SensitiveFloat{ParamType, NumType},
    sf_result::SensitiveFloat{ParamType, NumType},
    v::NumType, g_d::Vector{NumType}, g_h::Matrix{NumType};
    calculate_hessian::Bool=true)

  # TODO: time consuming **************

  # TODO: this line is allocating a lot of memory and I don't know why.
  # Commenting this line out attributes the same allocation to the next line.
  # Is memory being allocated lazily or misattributed?
  @assert g_h[1, 2] == g_h[2, 1]

  # You have to do this in the right order to not overwrite needed terms.<commit_msg>combine_sf_hessian! speedup: 1005 ms -> 456 ms<commit_after>
#using CelesteTypes.getids
#using CelesteTypes.HessianEntry

export multiply_sf!, add_scaled_sfs!, combine_sfs!, add_sources_sf!

@doc """
A function value and its derivative with respect to its arguments.

Attributes:
  v:  The value
  d:  The derivative with respect to each variable in
      P-dimensional VariationalParams for each of S celestial objects
      in a local_P x local_S matrix.
  h:  The second derivative with respect to each variational parameter,
      in the same format as d.  This is used for the full Hessian
      with respect to all the sources.
  hs: An array of per-source Hessians.  This will generally be reserved
      for the Hessian of brightness values that depend only on one source.
""" ->
type SensitiveFloat{ParamType <: CelesteTypes.ParamSet, NumType <: Number}
    v::NumType
    d::Matrix{NumType} # local_P x local_S
    # h is ordered so that p changes fastest.  For example, the indices
    # of a column of h correspond to the indices of d's stacked columns.
    h::Matrix{NumType} # (local_P * local_S) x (local_P * local_S)
    ids::ParamType
end


#########################################################

@doc """
Set a SensitiveFloat's hessian term, maintaining symmetry.
""" ->
function set_hess!{ParamType <: CelesteTypes.ParamSet, NumType <: Number}(
    sf::SensitiveFloat{ParamType, NumType},
    i::Int64, j::Int64, v::NumType)
  i != j ?
    sf.h[i, j] = sf.h[j, i] = v:
    sf.h[i, j] = v
end

function zero_sensitive_float{ParamType <: CelesteTypes.ParamSet}(
  ::Type{ParamType}, NumType::DataType, local_S::Int64)
    local_P = length(ParamType)
    d = zeros(NumType, local_P, local_S)
    h = zeros(NumType, local_P * local_S, local_P * local_S)
    SensitiveFloat{ParamType, NumType}(
      zero(NumType), d, h, getids(ParamType))
end

function zero_sensitive_float{ParamType <: CelesteTypes.ParamSet}(
  ::Type{ParamType}, NumType::DataType)
    # Default to a single source.
    zero_sensitive_float(ParamType, NumType, 1)
end

function clear!{ParamType <: CelesteTypes.ParamSet, NumType <: Number}(
  sp::SensitiveFloat{ParamType, NumType})
    sp.v = zero(NumType)
    fill!(sp.d, zero(NumType))
    fill!(sp.h, zero(NumType))
end

# If no type is specified, default to using Float64.
function zero_sensitive_float{ParamType <: CelesteTypes.ParamSet}(
  param_arg::Type{ParamType}, local_S::Int64)
    zero_sensitive_float(param_arg, Float64, local_S)
end

function zero_sensitive_float{ParamType <: CelesteTypes.ParamSet}(
  param_arg::Type{ParamType})
    zero_sensitive_float(param_arg, Float64, 1)
end

function +(sf1::SensitiveFloat, sf2::SensitiveFloat)
  S = size(sf1.d)[2]

  # Simply asserting equality of the ids doesn't work for some reason.
  @assert typeof(sf1.ids) == typeof(sf2.ids)
  @assert length(sf1.ids) == length(sf2.ids)
  [ @assert size(sf1.h[s]) == size(sf2.h[s]) for s=1:S ]

  @assert size(sf1.d) == size(sf2.d)

  sf3 = deepcopy(sf1)
  sf3.v = sf1.v + sf2.v
  sf3.d = sf1.d + sf2.d
  sf3.h = sf1.h + sf2.h

  sf3
end


@doc """
Updates sf1 in place with sf1 * sf2.

ids1 and ids2 are the ids for which sf1 and sf2 have nonzero derivatives,
respectively.
""" ->
function multiply_sf!{ParamType <: CelesteTypes.ParamSet, NumType <: Number}(
    sf1::SensitiveFloat{ParamType, NumType},
    sf2::SensitiveFloat{ParamType, NumType};
    ids1::Array{Int64}=collect(1:length(sf1.ids)),
    ids2::Array{Int64}=collect(1:length(sf2.ids)),
    calculate_hessian::Bool=true)

  S = size(sf1.d)[2]
  @assert S == 1 # For now this is only for the brightness calculations.
  # TODO: replace this with combine_sfs!.

  # You have to do this in the right order to not overwrite needed terms.

  if calculate_hessian
    # Second derivate terms involving second derivates.
    sf1.h[:, :] = sf1.v * sf2.h + sf2.v * sf1.h
    sf1.h[:, :] += sf1.d[:] * sf2.d[:]' + sf2.d[:] * sf1.d[:]'
  end
  sf1.d[:, :] = sf2.v * sf1.d + sf1.v * sf2.d

  sf1.v = sf1.v * sf2.v
end


@doc """
Factor out the hessian part of combine_sfs! to help the compiler.

TODO: I think this is a red herring and this can be put back in
""" ->
function combine_sfs_hessian!{ParamType <: CelesteTypes.ParamSet, NumType <: Number}(
    sf1::SensitiveFloat{ParamType, NumType},
    sf2::SensitiveFloat{ParamType, NumType},
    sf_result::SensitiveFloat{ParamType, NumType},
    g_d::Vector{NumType}, g_h::Matrix{NumType})

  const use_blas = false

  if use_blas
    # Chain rule for second derivatives.
    # BLAS for
    # sf_result.h[:, :] = g_d[1] * sf1.h + g_d[2] * sf2.h
    BLAS.blascopy!(prod(size(sf_result.h)), sf1.h, 1, sf_result.h, 1);
    BLAS.scal!(prod(size(sf_result.h)), g_d[1], sf_result.h, 1);
    BLAS.axpy!(g_d[2], sf2.h, sf_result.h)

    # BLAS for
    # sf_result.h[:, :] +=
    #   g_h[1, 1] * sf1.d[:] * sf1.d[:]' +
    #   g_h[2, 2] * sf2.d[:] * sf2.d[:]' +
    #   g_h[1, 2] * (sf1.d[:] * sf2.d[:]' + sf2.d[:] * sf1.d[:]')
    sf1d = sf1.d[:];
    sf2d = sf2.d[:];
    BLAS.ger!(g_h[1, 1], sf1d, sf1d, sf_result.h);
    BLAS.ger!(g_h[2, 2], sf2d, sf2d, sf_result.h);
    BLAS.ger!(g_h[1, 2], sf1d, sf2d, sf_result.h);
    BLAS.ger!(g_h[1, 2], sf2d, sf1d, sf_result.h);
  else
    p1, p2 = size(sf_result.h)
    for ind2 = 1:p2
      sf11_factor = g_h[1, 1] * sf1.d[ind2] + g_h[1, 2] * sf2.d[ind2]
      sf21_factor = g_h[1, 2] * sf1.d[ind2] + g_h[2, 2] * sf2.d[ind2]

      @inbounds for ind1 = 1:p1
        sf_result.h[ind1, ind2] = 
          g_d[1] * sf1.h[ind1, ind2] +
          g_d[2] * sf2.h[ind1, ind2] +
          sf11_factor * sf1.d[ind1] +
          sf21_factor * sf2.d[ind1]
      end
    end
  end
end


@doc """
Updates sf_result in place with g(sf1, sf2), where
g_d = (g_1, g_2) is the gradient of g and
g_h = (g_11, g_12; g_12, g_22) is the hessian of g,
each evaluated at (sf1, sf2).

The result is stored in sf_result.  The order is done in such a way that
it can overwrite sf1 or sf2 and still be accurate.
""" ->
function combine_sfs!{ParamType <: CelesteTypes.ParamSet, NumType <: Number}(
    sf1::SensitiveFloat{ParamType, NumType},
    sf2::SensitiveFloat{ParamType, NumType},
    sf_result::SensitiveFloat{ParamType, NumType},
    v::NumType, g_d::Vector{NumType}, g_h::Matrix{NumType};
    calculate_hessian::Bool=true)

  # TODO: time consuming **************

  # TODO: this line is allocating a lot of memory and I don't know why.
  # Commenting this line out attributes the same allocation to the next line.
  # Is memory being allocated lazily or misattributed?
  @assert g_h[1, 2] == g_h[2, 1]

  # You have to do this in the right order to not overwrite needed terms.
{'fraction_non_alphanumeric': 0.14748365516192793, 'fraction_numerical': 0.029876843545689525, 'mean_word_length': 3.7049356223175964, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4081596', 'n_tokens_mistral': 5188, 'n_tokens_neox': 4993, 'n_words': 1337}