text:
#pragma GCC diagnostic ignored "-Wdeclaration-after-statement"
/*
 * ocl_model_maintenance.c
 *
 *  Created on: 09.02.2014
 *      Author: mheimel
 */

#include "ocl_model_maintenance.h"

#include <math.h>
#include <stdlib.h>
#include <sys/time.h>

#include "ocl_adaptive_bandwidth.h"
#include "ocl_error_metrics.h"
#include "ocl_estimator.h"
#include "ocl_sample_maintenance.h"
#include "ocl_utilities.h"

// Optimization routines.
#include <nlopt.h>
#include "lbfgs/lbfgs.h"

#include "catalog/pg_kdefeedback.h"
#include "executor/spi.h"
#include "optimizer/path/gpukde/ocl_estimator_api.h"
#include "storage/lock.h"

// Global GUC variables
bool kde_enable_bandwidth_optimization;
int kde_bandwidth_optimization_feedback_window;
extern int kde_bandwidth_representation;

const double learning_rate = 0.01f;

// ############################################################
// # Code for adaptive bandwidth optimization (online learning).
// ############################################################

void ocl_notifyModelMaintenanceOfSelectivity(
    Oid relation, double selected, double allrows) {
  CREATE_TIMER();

  // Check if we have an estimator for this relation.
  ocl_estimator_t* estimator = ocl_getEstimator(relation);
  if (estimator == NULL) return;
  if (!estimator->open_estimation) return;  // No registered estimation.

  double selectivity = selected / allrows;
  estimator->rows_in_table = allrows;

  // Notify the sample maintenance of this observation so it can track the sample quality.
  ocl_notifySampleMaintenanceOfSelectivity(estimator, selectivity);

  // Update the bandwidth using online learning.
  ocl_runOnlineLearningStep(estimator, selectivity);

  // Write the error to the log file.
  ocl_reportErrorToLogFile(
      relation, estimator->last_selectivity, selectivity,
      estimator->rows_in_table);

  // We are done.
  estimator->open_estimation = false;
  LOG_TIMER("Model Maintenance");
}

// ############################################################
// # Code for offline bandwidth optimization (batch learning).
// ############################################################

// Helper function to extract the n latest feedback records for the given
// estimator from the catalogue. The function will only return tuples that
// have feedback that matches the estimator's attributes.
//
// Returns the actual number of valid feedback records in the catalog.
static unsigned int ocl_extractLatestFeedbackRecordsFromCatalog(
    ocl_estimator_t* estimator, unsigned int requested_records,
    kde_float_t* range_buffer, kde_float_t* selectivity_buffer) {
  unsigned int current_tuple = 0;

  // Open a new scan over the feedback table.
  unsigned int i, j;
  if (SPI_connect() != SPI_OK_CONNECT) {
    fprintf(stderr, "> Error connecting to Postgres Backend.\n");
    return current_tuple;
  }
  char query_buffer[1024];
  snprintf(query_buffer, 1024,
           "SELECT columns, ranges, alltuples, qualifiedtuples FROM "
           "pg_kdefeedback WHERE \"table\" = %i ORDER BY \"timestamp\" DESC "
           "LIMIT %i;", estimator->table, requested_records);
  if (SPI_execute(query_buffer, true, 0) != SPI_OK_SELECT) {
    fprintf(stderr, "> Error querying system table.\n");
    return current_tuple;
  }
  SPITupleTable *spi_tuptable = SPI_tuptable;
  unsigned int result_tuples = SPI_processed;
  TupleDesc spi_tupdesc = spi_tuptable->tupdesc;
  for (i = 0; i < result_tuples; ++i) {
    bool isnull;
    HeapTuple record_tuple = spi_tuptable->vals[i];
    // First, check whether this record only covers columns that are part of the estimator.
    unsigned int columns_in_record = DatumGetInt32(
          SPI_getbinval(record_tuple, spi_tupdesc, 1, &isnull));
    if ((columns_in_record | estimator->columns) != estimator->columns) continue;
    // This is a valid record, initialize the range buffer.
    unsigned int pos = current_tuple * 2 * estimator->nr_of_dimensions;
    for (j=0; j<estimator->nr_of_dimensions; ++j) {
      range_buffer[pos + 2*j] = -1.0 * INFINITY;
      range_buffer[pos + 2*j + 1] = INFINITY;
    }
    // Now extract all clauses and isert them into the range buffer.
    RQClause* clauses;
    unsigned int nr_of_clauses = extract_clauses_from_buffer(DatumGetByteaP(
        SPI_getbinval(record_tuple, spi_tupdesc, 2, &isnull)), &clauses);
    for (j=0; j<nr_of_clauses; ++j) {
      // First, locate the correct column position in the estimator.
      int column_in_estimator = estimator->column_order[clauses[j].var];
      // Re-Scale the bounds, add potential padding and write them to their position.
      float8 lo = clauses[j].lobound;
      if (clauses[j].loinclusive != EX) lo -= 0.001;
      float8 hi = clauses[j].hibound;
      if (clauses[j].hiinclusive != EX) hi += 0.001;
      range_buffer[pos + 2*column_in_estimator] = lo;
      range_buffer[pos + 2*column_in_estimator + 1] = hi;
    }
    // Finally, extract the selectivity and increase the tuple count.
    double all_rows = DatumGetFloat8(
        SPI_getbinval(record_tuple, spi_tupdesc, 3, &isnull));
    double qualified_rows = DatumGetFloat8(
        SPI_getbinval(record_tuple, spi_tupdesc, 4, &isnull));
    selectivity_buffer[current_tuple++] = qualified_rows / all_rows;
  }
  // We are done :)
  SPI_finish();
  return current_tuple;
}

// Helper function that extracts feedback for the given estimator and
// pushes it to the device.
static unsigned int ocl_prepareFeedback(
    ocl_estimator_t* estimator, cl_mem* device_ranges,
    cl_mem* device_selectivities) {
  cl_int err = CL_SUCCESS;
  // First, we have to count how many matching feedback records are available
  // for this estimator in the query feedback table.
  if (SPI_connect() != SPI_OK_CONNECT) {
    fprintf(stderr, "> Error connecting to Postgres Backend.\n");
    return 0;
  }
  char query_buffer[1024];
  snprintf(query_buffer, 1024,
           "SELECT COUNT(*) FROM pg_kdefeedback WHERE \"table\" = %i;",
           estimator->table);
  if (SPI_execute(query_buffer, true, 0) != SPI_OK_SELECT) {
    fprintf(stderr, "> Error querying system table.\n");
    return 0;
  }
  SPITupleTable *tuptable = SPI_tuptable;
  bool isnull;
  unsigned int available_records = DatumGetInt32(
      SPI_getbinval(tuptable->vals[0], tuptable->tupdesc, 1, &isnull));
  SPI_finish();

  if (available_records == 0) {
    fprintf(stderr, "> No feedback available for table %i\n", estimator->table);
    return 0;
  }

  // Adjust the number of records according to the specified window size.
  int used_records;
  if (kde_bandwidth_optimization_feedback_window == -1)
    used_records = available_records;
  else
    used_records = Min(available_records,
                       kde_bandwidth_optimization_feedback_window);
  fprintf(stderr, "> Checking the %i latest feedback records.\n", used_records);

  // Allocate arrays and fetch the actual feedback data.
  kde_float_t* range_buffer = palloc(
      sizeof(kde_float_t) * 2 * estimator->nr_of_dimensions * used_records);
  kde_float_t* selectivity_buffer = palloc(sizeof(kde_float_t) * used_records);
  unsigned int actual_records = ocl_extractLatestFeedbackRecordsFromCatalog(
      estimator, used_records, range_buffer, selectivity_buffer);

  // Now push the records to the device to prepare the actual optimization.
  if (actual_records == 0) {
    fprintf(stderr, "> No valid feedback records found.\n");
    return 0;
  }
  fprintf(stderr, "> Found %i valid records, pushing to device.\n",
          actual_records);
  ocl_context_t* context = ocl_getContext();
  *device_ranges = clCreateBuffer(
      context->context, CL_MEM_READ_WRITE,
      sizeof(kde_float_t) * 2 * actual_records * estimator->nr_of_dimensions,
      NULL, &err);
  Assert(err == CL_SUCCESS);

  err = clEnqueueWriteBuffer(
      context->queue, *device_ranges, CL_FALSE, 0,
      sizeof(kde_float_t) * 2 * actual_records * estimator->nr_of_dimensions,
      range_buffer, 0, NULL, NULL);
  estimator->stats->optimization_transfer_to_device++;
  Assert(err == CL_SUCCESS);

  *device_selectivities = clCreateBuffer(
      context->context, CL_MEM_READ_WRITE,
      sizeof(kde_float_t) * actual_records, NULL, &err);
  Assert(err == CL_SUCCESS);

  err = clEnqueueWriteBuffer(
      context->queue, *device_selectivities, CL_FALSE, 0,
      sizeof(kde_float_t) * actual_records, selectivity_buffer, 0, NULL, NULL);
  estimator->stats->optimization_transfer_to_device++;
  Assert(err == CL_SUCCESS);

  err = clFinish(context->queue);
  Assert(err == CL_SUCCESS);

  pfree(range_buffer);
  pfree(selectivity_buffer);
  return actual_records;
}

typedef struct {
  ocl_estimator_t* estimator;
  unsigned int nr_of_observations;
  cl_mem observed_ranges;
  cl_mem observed_selectivities;
  // Temporary buffers.
  size_t stride_size;
  cl_mem gradient_accumulator_buffer;
  cl_mem error_accumulator_buffer;
  cl_mem gradient_buffer;
  cl_mem error_buffer;
  // Required event and accumulation buffers.
  ocl_aggregation_descriptor_t** summation_descriptors;
  cl_mem* summation_buffers;
} optimization_config_t;

/**
 * Counters for evaluations.
 */
int evaluations;
double start_error;
struct timeval opt_start;

/*
 * Function to compute the gradient for a penalized objective function that will
 * add a strong penalty factor to negative bandwidth values. This function is
 * passed to nlopt to compute the gradient and evaluate the function.
 */

/*
 * Callback function that computes the gradient and value for the objective
 * function at the current bandwidth.
 */
static double computeGradient(
    unsigned n, const double* bandwidth, double* gradient, void* params) {
  unsigned int i;
  cl_int err = CL_SUCCESS;
  optimization_config_t* conf = (optimization_config_t*)params;
  ocl_estimator_t* estimator = conf->estimator;
  ocl_context_t* context = ocl_getContext();

  evaluations++;

  if (ocl_isDebug()) {
    fprintf(stderr, ">>> Evaluation %i:\n\tCurrent bandwidth:", evaluations);
    for (i=0; i<n; ++i) fprintf(stderr, " %e", bandwidth[i]);
    fprintf(stderr, "\n");
  }

  // First, transfer the current bandwidth to the device. Note that we might
  // need to cast the bandwidth to float first.
  kde_float_t* fbandwidth = NULL;
  cl_event input_transfer_event;
  if (sizeof(kde_float_t) != sizeof(double)) {
    fbandwidth = palloc(sizeof(kde_float_t) * estimator->nr_of_dimensions);
    for (i = 0; i<estimator->nr_of_dimensions; ++i) {
      fbandwidth[i] = bandwidth[i];
    }
    err = clEnqueueWriteBuffer(
        context->queue, estimator->bandwidth_buffer, CL_FALSE,
        0, sizeof(kde_float_t) * estimator->nr_of_dimensions,
        fbandwidth, 0, NULL, &input_transfer_event);
    estimator->stats->optimization_transfer_to_device++;
    Assert(err == CL_SUCCESS);
  } else {
    err = clEnqueueWriteBuffer(
        context->queue, estimator->bandwidth_buffer, CL_FALSE,
        0, sizeof(kde_float_t) * estimator->nr_of_dimensions,
        bandwidth, 0, NULL, &input_transfer_event);
    estimator->stats->optimization_transfer_to_device++;
    Assert(err == CL_SUCCESS);
  }
  // Prepare the kernel that computes a gradient for each observation.
  cl_kernel gradient_kernel = ocl_getKernel(
      ocl_getSelectedErrorMetric()->batch_kernel_name,
      estimator->nr_of_dimensions);
    // First, fix the local and global size. We identify the optimal local size
    // by looking at available and required local memory and by ensuring that
    // the local size is evenly divisible by the preferred workgroup multiple..
  size_t local_size;
  err = clGetKernelWorkGroupInfo(
      gradient_kernel, context->device, CL_KERNEL_WORK_GROUP_SIZE,
      sizeof(size_t), &local_size, NULL);
  Assert(err == CL_SUCCESS);

  size_t available_local_memory;
  err = clGetKernelWorkGroupInfo(
      gradient_kernel, context->device, CL_KERNEL_LOCAL_MEM_SIZE,
      sizeof(size_t), &available_local_memory, NULL);
  Assert(err == CL_SUCCESS);

  available_local_memory = context->local_mem_size - available_local_memory;
  local_size = Min(
      local_size,
      available_local_memory / (3 * sizeof(kde_float_t) * estimator->nr_of_dimensions));
  size_t preferred_local_size_multiple;
  err = clGetKernelWorkGroupInfo(
      gradient_kernel, context->device,
      CL_KERNEL_PREFERRED_WORK_GROUP_SIZE_MULTIPLE,
      sizeof(size_t), &preferred_local_size_multiple, NULL);
  Assert(err == CL_SUCCESS);

  local_size = preferred_local_size_multiple
      * (local_size / preferred_local_size_multiple);
  size_t global_size = local_size * (conf->nr_of_observations / local_size);
  if (global_size < conf->nr_of_observations) global_size += local_size;
    // Configure the kernel by setting all required parameters.
  err |= clSetKernelArg(
      gradient_kernel, 0, sizeof(cl_mem), &(estimator->sample_buffer));
  err |= clSetKernelArg(
      gradient_kernel, 1, sizeof(unsigned int), &(estimator->rows_in_sample));
  err |= clSetKernelArg(
      gradient_kernel, 2, sizeof(cl_mem), &(conf->observed_ranges));
  err |= clSetKernelArg(
      gradient_kernel, 3, sizeof(cl_mem), &(conf->observed_selectivities));
  err |= clSetKernelArg(
      gradient_kernel, 4, sizeof(unsigned int), &(conf->nr_of_observations));
  err |= clSetKernelArg(
      gradient_kernel, 5, sizeof(cl_mem), &(estimator->bandwidth_buffer));
  err |= clSetKernelArg(
      gradient_kernel, 6,
      local_size * sizeof(kde_float_t) * estimator->nr_of_dimensions, NULL);
  err |= clSetKernelArg(
      gradient_kernel, 7,
      local_size * sizeof(kde_float_t) * estimator->nr_of_dimensions, NULL);
  err |= clSetKernelArg(
      gradient_kernel, 8,
      local_size * sizeof(kde_float_t) * estimator->nr_of_dimensions, NULL);
  err |= clSetKernelArg(
      gradient_kernel, 9, sizeof(cl_mem), &(conf->error_accumulator_buffer));
  err |= clSetKernelArg(
      gradient_kernel, 10, sizeof(cl_mem),
      &(conf->gradient_accumulator_buffer));
  unsigned int stride_elements = conf->stride_size / sizeof(kde_float_t);
  err |= clSetKernelArg(
      gradient_kernel, 11, sizeof(unsigned int), &stride_elements);
  err |= clSetKernelArg(
      gradient_kernel, 12, sizeof(unsigned int), &(estimator->rows_in_table));
  err |= clSetKernelArg(
      gradient_kernel, 13, sizeof(cl_mem), &(estimator->mean_buffer));
  err |= clSetKernelArg(
      gradient_kernel, 14, sizeof(cl_mem), &(estimator->sdev_buffer));
  Assert(err == CL_SUCCESS);

  // Compute the gradient for each observation.
  cl_event partial_gradient_event;
  err = clEnqueueNDRangeKernel(
      context->queue, gradient_kernel, 1, NULL, &global_size, &local_size, 1,
      &input_transfer_event, &partial_gradient_event);
  Assert(err == CL_SUCCESS);

  // Sum up the individual error contributions ...
  cl_event* summation_events = palloc(
      sizeof(cl_event) * (1 + estimator->nr_of_dimensions));

  summation_events[0] = predefinedSumOfArray(
      conf->summation_descriptors[0], partial_gradient_event);
  // .. and the individual gradients.
  for (i=0; i<estimator->nr_of_dimensions; ++i) {
    summation_events[i + 1] = predefinedSumOfArray(
        conf->summation_descriptors[i + 1], partial_gradient_event);
  }
  // Now transfer the gradient back to the device.
  cl_event result_events[2];
  kde_float_t* tmp_gradient = palloc(
      sizeof(kde_float_t) * estimator->nr_of_dimensions);
  err = clEnqueueReadBuffer(
      context->queue, conf->gradient_buffer, CL_FALSE,
      0, sizeof(kde_float_t) * estimator->nr_of_dimensions,
      tmp_gradient, estimator->nr_of_dimensions + 1,
      summation_events, &(result_events[0]));
  estimator->stats->optimization_transfer_to_host++;
  Assert(err == CL_SUCCESS);

  // As well as the error.
  kde_float_t error;
  err = clEnqueueReadBuffer(
      context->queue, conf->error_buffer, CL_FALSE,
      0, sizeof(kde_float_t), &error,
      1, &(summation_events[0]),  &(result_events[1]));
  estimator->stats->optimization_transfer_to_host++;
  Assert(err == CL_SUCCESS);

  err = clWaitForEvents(2, result_events);
  Assert(err == CL_SUCCESS);

  if (err != 0) {
    fprintf(stderr, "OpenCL functions failed to compute gradient.\n");
  }

  error /= conf->nr_of_observations;
  if (evaluations == 1) start_error = error;
  struct timeval now; gettimeofday(&now, NULL);
  long seconds = now.tv_sec - opt_start.tv_sec;
  long useconds = now.tv_usec - opt_start.tv_usec;
  long mtime = ((seconds) * 1000 + useconds / 1000.0) + 0.5;
  // Ok, cool. Transfer the bandwidth back.
  if (ocl_isDebug()) {
    fprintf(
        stderr, "\r\tOptimization round %i. Current error: %f "
        "(started at %f), took: %ld ms.",
        evaluations, error, start_error, mtime);
  }
  // Finally, cast back to double.
  if (gradient) {
    for (i = 0; i<estimator->nr_of_dimensions; ++i) {
    // Apply the gradient normalization.
      double h = bandwidth[i];
      if(kde_bandwidth_representation == PLAIN_BW){
          gradient[i] = tmp_gradient[i] * M_SQRT2 / (
              sqrt(M_PI) * h * h * pow(2.0, estimator->nr_of_dimensions) *
              conf->nr_of_observations * estimator->rows_in_sample);
      }
      else {
          gradient[i] = tmp_gradient[i] * M_SQRT2 / (
              sqrt(M_PI) * exp(h) * pow(2.0, estimator->nr_of_dimensions) *
              conf->nr_of_observations * estimator->rows_in_sample);
      }
    }
    if (ocl_isDebug()) {
      fprintf(stderr, "\n\tGradient:");
      for (i=0; i<n; ++i) fprintf(stderr, " %e", gradient[i]);
      fprintf(stderr, "\n");
    }
  }
  // Ok, clean everything up.
  for (i=0; i<estimator->nr_of_dimensions; ++i) {
    err = clReleaseEvent(summation_events[i]);
    Assert(err == CL_SUCCESS);
  }
  pfree(summation_events);
  pfree(tmp_gradient);
  if (fbandwidth) pfree(fbandwidth);
  err |= clReleaseEvent(input_transfer_event);
  err |= clReleaseEvent(partial_gradient_event);
  err |= clReleaseEvent(result_events[0]);
  err |= clReleaseEvent(result_events[1]);
  err |= clReleaseKernel(gradient_kernel);
  Assert(err == CL_SUCCESS);

  return error;
}

static void ocl_setScottsBandwidth(ocl_estimator_t* estimator) {
  ocl_context_t* context = ocl_getContext();
  cl_int err = CL_SUCCESS;
  // First, we need to compute the variance for each dimension.
  unsigned int i;
  kde_float_t* variances = malloc(
      sizeof(kde_float_t) * estimator->nr_of_dimensions);
  cl_mem* buffers = malloc(sizeof(cl_mem) * estimator->nr_of_dimensions);
  cl_mem averages = clCreateBuffer(
      context->context, CL_MEM_READ_WRITE,
      sizeof(kde_float_t) * estimator->nr_of_dimensions, NULL, &err);
  Assert(err == CL_SUCCESS);
  cl_event* events = malloc(sizeof(cl_event) * estimator->nr_of_dimensions);
  size_t sample_size = estimator->rows_in_sample;
  size_t dimensions = estimator->nr_of_dimensions;
  for (i=0; i<estimator->nr_of_dimensions; ++i) {
    // Allocate all required buffers.
    buffers[i] = clCreateBuffer(
        context->context, CL_MEM_READ_WRITE,
        sizeof(kde_float_t) * estimator->rows_in_sample, NULL, &err);
    Assert(err == CL_SUCCESS);
    // First we extract all sample components for this dimension.
    cl_kernel extractComponents = ocl_getKernel("extract_dimension", 0);
    err |= clSetKernelArg(
        extractComponents, 0, sizeof(cl_mem), &(estimator->sample_buffer));
    err |= clSetKernelArg(extractComponents, 1, sizeof(cl_mem), &(buffers[i]));
    err |= clSetKernelArg(extractComponents, 2, sizeof(unsigned int), &i);
    err |= clSetKernelArg(extractComponents, 3, sizeof(unsigned int),
        &(estimator->nr_of_dimensions));
    Assert(err == CL_SUCCESS);

    cl_event extraction_event;
    err = clEnqueueNDRangeKernel(
        context->queue, extractComponents, 1, NULL, &sample_size, NULL,
        0, NULL, &extraction_event);
    Assert(err == CL_SUCCESS);
    // Now we sum them up, so we can compute the average.
    cl_event average_summation_event = sumOfArray(
        buffers[i], estimator->rows_in_sample, averages, i, extraction_event);
    // Alright, we can compute the variance contributions from each point.
    cl_kernel precomputeVariance = ocl_getKernel("precompute_variance", 0);
    err |= clSetKernelArg(precomputeVariance, 0, sizeof(cl_mem), &(buffers[i]));
    err |= clSetKernelArg(precomputeVariance, 1, sizeof(cl_mem), &averages);
    err |= clSetKernelArg(precomputeVariance, 2, sizeof(unsigned int), &i);
    err |= clSetKernelArg(precomputeVariance, 3, sizeof(unsigned int),
        &(estimator->rows_in_sample));
    cl_event variance_event;
    err = clEnqueueNDRangeKernel(
        context->queue, precomputeVariance, 1, NULL, &sample_size, NULL,
        1, &average_summation_event, &variance_event);
    Assert(err == CL_SUCCESS);

    // We now sum up the single contributions to compute the variance.
    cl_event variance_summation_event = sumOfArray(
        buffers[i], estimator->rows_in_sample, averages, i, variance_event);
    // Finally, we can compute and store the bandwidth for this value.
    cl_kernel finalizeBandwidth = ocl_getKernel("set_scotts_bandwidth", 0);
    err |= clSetKernelArg(finalizeBandwidth, 0, sizeof(cl_mem), &averages);
    err |= clSetKernelArg(finalizeBandwidth, 1, sizeof(cl_mem),
        &(estimator->bandwidth_buffer));
    err |= clSetKernelArg(finalizeBandwidth, 2, sizeof(unsigned int), &i);
    err |= clSetKernelArg(finalizeBandwidth, 3, sizeof(unsigned int),
        &(estimator->nr_of_dimensions));
    err |= clSetKernelArg(finalizeBandwidth, 4, sizeof(unsigned int),
        &(estimator->rows_in_sample));
    Assert(err == CL_SUCCESS);

    err = clEnqueueNDRangeKernel(
        context->queue, finalizeBandwidth, 1, NULL, &dimensions, NULL,
        1, &variance_summation_event, &events[i]);
    Assert(err == CL_SUCCESS);
    // Clean up.
    err = clReleaseMemObject(buffers[i]);
    Assert(err == CL_SUCCESS);

    err |= clReleaseEvent(extraction_event);
    err |= clReleaseEvent(average_summation_event);
    err |= clReleaseEvent(variance_event);
    err |= clReleaseEvent(variance_summation_event);
    Assert(err == CL_SUCCESS);
  }
  err = clReleaseMemObject(averages);
  Assert(err == CL_SUCCESS);
  // Wait for the events to finalize:
  err = clWaitForEvents(estimator->nr_of_dimensions, events);
  Assert(err == CL_SUCCESS);
  for (i=0; i<estimator->nr_of_dimensions; ++i) {
    err = clReleaseEvent(events[i]);
    Assert(err == CL_SUCCESS);
  }
  // Clean up.
  free(variances);
  free(events);
}

void ocl_runModelOptimization(ocl_estimator_t* estimator) {
  if (estimator == NULL) return;
  cl_int err = CL_SUCCESS;
  // Set the rule-of-thumb bandwidth to initialize the estimator.
  ocl_setScottsBandwidth(estimator);
  // Now check if we do a full bandwidth optimization.
  if (!kde_enable_bandwidth_optimization) return;
  if (ocl_isDebug()) {
    fprintf(
        stderr, "Beginning model optimization for estimator on table %i\n",
        estimator->table);
  }
  // First, we need to fetch the feedback records for this table and push them
  // to the device.
  cl_mem device_ranges, device_selectivites;
  unsigned int feedback_records = ocl_prepareFeedback(
      estimator, &device_ranges, &device_selectivites);
  if (feedback_records == 0) return;

  // We need to transfer the bandwidth to the host.
  kde_float_t* fbandwidth = palloc(
      sizeof(kde_float_t) * estimator->nr_of_dimensions);
  ocl_context_t* context = ocl_getContext();
  err = clEnqueueReadBuffer(
      context->queue, estimator->bandwidth_buffer, CL_TRUE, 0,
      sizeof(kde_float_t) * estimator->nr_of_dimensions,
      fbandwidth, 0, NULL, NULL);
  estimator->stats->optimization_transfer_to_host++;
  Assert(err == CL_SUCCESS);
  // Cast to double (nlopt operates on double).
  double* bandwidth = lbfgs_malloc(estimator->nr_of_dimensions);
  unsigned int i;
  for (i=0; i<estimator->nr_of_dimensions; ++i) {
    bandwidth[i] = fbandwidth[i];
  }
  // Package all required buffers.
  optimization_config_t params;
  params.estimator = estimator;
  params.nr_of_observations = feedback_records;
  params.observed_ranges = device_ranges;
  params.observed_selectivities = device_selectivites;
  params.error_accumulator_buffer = clCreateBuffer(
      context->context, CL_MEM_READ_WRITE,
      sizeof(kde_float_t) * feedback_records, NULL, &err);
  Assert(err == CL_SUCCESS);
  // Allocate a buffer to hold temporary gradient contributions. This buffer
  // will keep D contributions per observation. We store all contributions
  // consecutively (i.e. 111222333444). For optimal performance, we therefore
  // have to make sure that the consecutive regions (strides) have a size that
  // is aligned to the required machine alignment.
  params.stride_size = sizeof(kde_float_t) * feedback_records;
  if ((params.stride_size * 8) % context->required_mem_alignment) {
    // The stride size is misaligned, add some padding.
    params.stride_size *= 8;
    params.stride_size =
        (1 + params.stride_size / context->required_mem_alignment)
        * context->required_mem_alignment;
    params.stride_size /= 8;
  }
  params.gradient_accumulator_buffer = clCreateBuffer(
      context->context, CL_MEM_READ_WRITE,
      estimator->nr_of_dimensions * params.stride_size,
      NULL, &err);
  Assert(err == CL_SUCCESS);
  params.gradient_buffer = clCreateBuffer(
      context->context, CL_MEM_READ_WRITE,
      estimator->nr_of_dimensions * sizeof(kde_float_t), NULL, &err);
  Assert(err == CL_SUCCESS);
  params.error_buffer = clCreateBuffer(
      context->context, CL_MEM_READ_WRITE, sizeof(kde_float_t), NULL, &err);
  Assert(err == CL_SUCCESS);
  // Prepare the summation buffers.
  params.summation_descriptors = palloc(
      sizeof(ocl_aggregation_descriptor_t) * (1 + estimator->nr_of_dimensions));
  // Prepare the partial aggregations for the error.
  params.summation_descriptors[0] = prepareSumDescriptor(
      params.error_accumulator_buffer, params.nr_of_observations,
      params.error_buffer, 0);
  // .. and for each dimension.
  params.summation_buffers = palloc(
      sizeof(cl_mem) * estimator->nr_of_dimensions);
  for (i=0; i<estimator->nr_of_dimensions; ++i) {
    cl_buffer_region region;
    region.size = params.stride_size;
    region.origin = i * params.stride_size;
    params.summation_buffers[i] = clCreateSubBuffer(
        params.gradient_accumulator_buffer, CL_MEM_READ_ONLY,
        CL_BUFFER_CREATE_TYPE_REGION, &region, &err);
    Assert(err == CL_SUCCESS);
    params.summation_descriptors[i + 1] = prepareSumDescriptor(
        params.summation_buffers[i], params.nr_of_observations,
        params.gradient_buffer, i);
  }

  // Ok, we are prepared. Call the optimization routine.
  gettimeofday(&opt_start, NULL);
  evaluations = 0;
  fprintf(stderr, "> Starting numerical optimization of the model.\n");
  // Prepare the bound constraints.
  double* lower_bounds = palloc(sizeof(double) * estimator->nr_of_dimensions);
  double* upper_bounds = palloc(sizeof(double) * estimator->nr_of_dimensions);
  if(kde_bandwidth_representation == LOG_BW){
    for (i=0; i<estimator->nr_of_dimensions; ++i) {
      lower_bounds[i] = log(1e-20);    // We never want to be negative.
    }
    for (i=0; i<estimator->nr_of_dimensions; ++i) {
      upper_bounds[i] = log(4) + bandwidth[i];    // We never want to be negative.
    }
  }
  else {
    for (i=0; i<estimator->nr_of_dimensions; ++i) {
      lower_bounds[i] = 1e-20;    // We never want to be negative.
    }
    for (i=0; i<estimator->nr_of_dimensions; ++i) {
      upper_bounds[i] = 4 * bandwidth[i];
    }
  }
  double tmp;
  int nl_err;
  // We start with a global optimization step.
  nlopt_opt global_optimizer = nlopt_create(
      NLOPT_GD_MLSL, estimator->nr_of_dimensions);
  Assert(global_optimizer);
  nl_err = nlopt_set_lower_bounds(global_optimizer, lower_bounds);
  Assert(nl_err > 0 );
  nl_err = nlopt_set_upper_bounds(global_optimizer, upper_bounds);
  Assert(nl_err > 0 );
  nl_err = nlopt_set_maxeval(global_optimizer, 120);
  Assert(nl_err > 0 );
  nl_err = nlopt_set_min_objective(global_optimizer, computeGradient, &params);
  Assert(nl_err > 0 );
  // Register a local LBFGS instance for the global optimizer.
  nlopt_opt global_local_optimizer = nlopt_create(
      NLOPT_LD_LBFGS, estimator->nr_of_dimensions);
  nl_err = nlopt_set_maxeval(global_local_optimizer, 40);
  Assert(nl_err > 0 );
  nl_err = nlopt_set_local_optimizer(global_optimizer, global_local_optimizer);
  Assert(nl_err > 0 );
  fprintf(stderr, "> Running global pre-optimization: ");
  nl_err = nlopt_optimize(global_optimizer, bandwidth, &tmp);
  Assert(nl_err > 0 );
  fprintf(stderr, " done (%i, %f)\n", nl_err, tmp);
  // Prepare the local refinement.
  nlopt_opt local_optimizer = nlopt_create(
      NLOPT_LD_LBFGS, estimator->nr_of_dimensions);
  nlopt_set_lower_bounds(local_optimizer, lower_bounds);
  nlopt_set_min_objective(local_optimizer, computeGradient, &params);
  nlopt_set_ftol_abs(local_optimizer, 1e-20);
  nlopt_set_maxeval(local_optimizer, 100);
  fprintf(stderr, "> Running local refinement: ");
  err = nlopt_optimize(local_optimizer, bandwidth, &tmp);
  fprintf(stderr, " done (%i, %f)\n", err, tmp);
  if (ocl_isDebug()) {
    if (err < 0) {
      fprintf(stderr, "\nOptimization failed: %i!", err);
    } else {
      fprintf(stderr, "\nNew bandwidth:");
      for ( i = 0; i < estimator->nr_of_dimensions ; ++i)
        fprintf(stderr, " %e", bandwidth[i]);
      fprintf(stderr, "\n");
    }
  }
  nlopt_destroy(local_optimizer);
  // Transfer the bandwidth to the device.
  for (i=0; i<estimator->nr_of_dimensions; ++i) {
    fbandwidth[i] = bandwidth[i];
  }
  err = clEnqueueWriteBuffer(
      context->queue, estimator->bandwidth_buffer, CL_TRUE, 0,
      sizeof(kde_float_t) * estimator->nr_of_dimensions,
      fbandwidth, 0, NULL, NULL);
  estimator->stats->optimization_transfer_to_device++;
  Assert(err == CL_SUCCESS);
  // Clean up.
  pfree(fbandwidth);
  lbfgs_free(bandwidth);
  for (i=0; i<estimator->nr_of_dimensions; ++i) {
    err = clReleaseMemObject(params.summation_buffers[i]);
    Assert(err == CL_SUCCESS);
    releaseAggregationDescriptor(params.summation_descriptors[i + 1]);
  }
  releaseAggregationDescriptor(params.summation_descriptors[0]);
  pfree(params.summation_descriptors);
  pfree(params.summation_buffers);
  err |= clReleaseMemObject(params.error_buffer);
  err |= clReleaseMemObject(params.gradient_buffer);
  err |= clReleaseMemObject(params.error_accumulator_buffer);
  err |= clReleaseMemObject(params.gradient_accumulator_buffer);
  err |= clReleaseMemObject(device_ranges);
  err |= clReleaseMemObject(device_selectivites);
  Assert(err == CL_SUCCESS);
}

max_stars_repo_name: sfu-db/feedback-kde

max_issues_repo_name: sfu-db/feedback-kde

max_forks_repo_name: sfu-db/feedback-kde


text:
#ifndef SOURCE_ROBANALYT_H_
#define SOURCE_ROBANALYT_H_

// *** source/rob_analyt.h ***
// Author: Kevin Wolz, date: 08/2018
//
// Class RobAnalyt: analytical algorithm to calculate distribution in the WPL.

#include <iostream>
#include <fstream>
#include <string>
#include <sstream>
#include <vector>
#include <stdlib.h>
#include <sys/stat.h>

#include <gsl/gsl_math.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_cblas.h>
#include <gsl/gsl_eigen.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_integration.h>

#include "utils.h"
#include "print_to_file.h"

class RobAnalyt {
 public:
 	gsl_matrix* gaussian_quadratic_matrix_;		// quadratic contraction with data
	unsigned int num_steps_;		 							// number of PDF evaluation points
	double lower_limit_;				 							// lower limit of PDF calculation
	double upper_limit_;				 							// upper limit of PDF calculation

	RobAnalyt(int num_steps,
	 					double upper_limit,
	 					double lower_limit,
	          unsigned int datdim, 							// number of data points
	          unsigned int pardim, 							// number of models parameters
	          gsl_matrix* g_matrix, 						// model design matrix
	          gsl_matrix* inv_covar_matrix, 		// inverse data covariance matrix
	          gsl_matrix* fisher_matrix_tot, 		// Fisher matrix of total data
						gsl_matrix* fisher_matrix_1, 			// Fisher matrix of first subset
						gsl_matrix* fisher_matrix_2, 			// Fisher matrix of second subset
						gsl_matrix* prior_matrix,
						gsl_vector* prior_param_values, 	// prior fiducial parameter values
						gsl_vector* data_mean_values); 		// mean of data vector

	int CalculateAnalytWeakPrior(std::string& filepath);

 private:
	unsigned int datdim_;
	unsigned int pardim_;
	gsl_matrix* g_matrix_tot_;
	gsl_matrix* g_matrix_1_;
	gsl_matrix* g_matrix_2_;
	gsl_matrix* inv_covar_matrix_tot_;
	gsl_matrix* inv_covar_matrix_1_;
	gsl_matrix* inv_covar_matrix_2_;
	gsl_matrix* fisher_matrix_tot_;
	gsl_matrix* fisher_matrix_1_;
	gsl_matrix* fisher_matrix_2_;
	gsl_matrix* prior_matrix_;
	gsl_vector* prior_param_values_;
	gsl_vector* data_mean_values_;

	double delta_y_; 											// constant shift of Robustness (1)
	double delta_y_wpl_;									// constant shift of Robustness (WPL, 1)
	double rob_constant_;									// constant shift of Robustness (2)
	gsl_vector* linear_gaussian_vector_; 	// linear contraction with data
	gsl_matrix* covar_matrix_tot_; 				// covariance matrix of total data
};

#endif // SOURCE_ROBANALYT_H_

ext: h

max_stars_repo_name: kevguitar/Robustness

max_issues_repo_name: kevguitar/Robustness

max_forks_repo_name: kevguitar/Robustness


text:
#include "lah.h"

#ifdef HAVE_LAPACK /* Use a LAPACK package to do the heavy work */
#include <lapacke.h>
lah_Return lah_chol(lah_mat *P, lah_index setZero)
{
    lah_index i, j;
    lapack_int res;
    if (P == NULL || !LAH_ISTYPE(P, lahTypeSquare))
        return lahReturnParameterError;

    res = POTRF(LAH_LAPACK_LAYOUT, 'L' , P->nR, P->data, P->nR);

    /* Set rest of matrix to zero, if setZero == 1 */
    if (setZero)
    {
        for (i = 0; i < P->nR; i++)
        {
            for(j = i + 1; j < P->nC; ++j)
            {
                LAH_ENTRY(P, i, j) = 0.0;
            }
        }
    }
    return (res == 0) ? lahReturnOk : lahReturnExternError;
}

#else /* fallback to primitive linear algebra subroutines */
#include <math.h>

lah_Return lah_chol(lah_mat *P, lah_index setZero)
{
    lah_index i, j, k, N;
    lah_value *diag, *value;

    if (P == NULL)
    {
        return lahReturnParameterError;
    }

    N = P->nR;

    /* see https://de.wikipedia.org/wiki/Cholesky-Zerlegung#Pseudocode */
    for (i = 0; i < N; i++)
    {
        /* Pointer to diagonal element */
        diag = P->data + i * (P->incRow + P->incCol);
        /* check if positive */
        if (*diag < 0.00001 )
        {
            return lahReturnMathError;
        }
        /* square root of diagonal element in place */
        *diag = sqrt(*diag);

        /* divide elements below diagonal by diagonal element */
        for (k = 1; k < N - i; k++)
        {
            *(diag + k * P->incRow) /= *diag;
        }
        /* compute upper right submatrix */
        for (j = i + 1; j < N; j++)
        {
            value = P->data + j * P->incRow + j * P->incCol;
            for (k = j; k < N; k++)
            {
                *value -= LAH_ENTRY(P, k, i) * LAH_ENTRY(P, j, i);
                value += P->incRow;
            }
        }
    }
    /* Set rest of matrix to zero, if setZero == 1 */
    if (setZero)
    {
        for (i = 0; i < N; i++)
        {
            for(j = i + 1; j < N; ++j)
            {
                LAH_ENTRY(P, i, j) = 0.0;
            }
        }
    }

    return lahReturnOk;
}
#endif /* endif primitive implementation */

ext: c


text:
#include <gsl/gsl_math.h>
#include "gsl_cblas.h"
#include "cblas.h"

void
cblas_drotg (double *a, double *b, double *c, double *s)
{
#define BASE double
#include "source_rotg.h"
#undef BASE
}


text:
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include <assert.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_spline.h>

#include "core_allvars.h"
#include "model_dust.h"
#include "model_misc.h"

double integrate_arr(const double arr1[MAX_STRING_LEN], const double arr2[MAX_STRING_LEN], const int npts, const double lower_limit, const double upper_limit);
double interpolate_arr(const double arr1[MAX_STRING_LEN], const double arr2[MAX_STRING_LEN], const int npts, const int xi);
double compute_imf (const double m);
double compute_taum (const double m);

void produce_metals_dust(const double metallicity, const double dt, const int p, const int centralgal, const int step, struct GALAXY *galaxies)
{
  double A = run_params.BinaryFraction;

  double yield;
  double yCr_snia, yFe_snia, yNi_snia;
  double yC_agb, yN_agb, yO_agb;
  double yC_sn, yO_sn, yMg_sn, ySi_sn, yS_sn, yCa_sn, yFe_sn;
  double Cr_snia, Fe_snia, Ni_snia;
  double C_agb, N_agb, O_agb;
  double C_sn, O_sn, Mg_sn, Si_sn, S_sn, Ca_sn, Fe_sn;
  double yCagb[run_params.countagb], yNagb[run_params.countagb], yOagb[run_params.countagb];
  double m_agb[run_params.countagb];
  double yCsn[run_params.countsn], yOsn[run_params.countsn], yMgsn[run_params.countsn], ySisn[run_params.countsn], ySsn[run_params.countsn], yCasn[run_params.countsn], yFesn[run_params.countsn];
  double m_sn[run_params.countsn];

  int i, j;
  double phi, taum, time, sfr;
  double gamma = 2.0; //to compute binary mass function for SN Ia
  double low_agb = 1; //lower limit for stars that ends up as AGB, Msun
  double low_binary = 3; //lower limit for binary star, Msun
  double up_agb = 6; //upper limit for stars that ends up as AGB, Msun
  double low_sn = run_params.msn[0]; //lower limit for stars that ends up as SN II, Msun
  double up_binary = 16; //upper limit for binary stars, Msun
  double up_sn = 40; //upper limit for stars that ends up as SN II, Msun
  double age[SNAPLEN], sfh[SNAPLEN];

  for (i=0; i<SNAPLEN; i++)
  {
        age[i] = run_params.lbtime[i]; //change the variable name to represent the age of universe instead of lbtime
        sfh[i] = galaxies[p].Sfr[i];
  }

/* the following computation is from eq. 9 Arrigoni et al. 2010 */

  // SN Ia //
  int count = 20; //any nymber will do
  double mu[count]; // M2/Mbinary
  double mbin[count]; //total mass of binary
  double fmu, yCrmu[count], yFemu[count], yNimu[count]; //to integrate over mu
  double yCrphi[count], yFephi[count], yNiphi[count];
  double max_mu = 0.5; //maximum value of mu

  for (i=0; i<count; i++){
        mbin[i] = low_binary + ((up_binary - low_binary) / (count-i));
        mu[i] = max_mu/(count-i);
        fmu = pow(2, 1+gamma) * (1+gamma) * pow(mu[i], gamma); //eq. 10 Arrigoni et al. 2010
        taum = compute_taum(mu[i] * mbin[i]);
        time = age[galaxies[p].SnapNum] - taum;
        if(time < run_params.lbtime[0]){
                sfr=0;
        }
        else {
                sfr = interpolate_arr(age, sfh, SNAPLEN, time);
        }
        yCrmu[i] = fmu * sfr * run_params.qCrsnia;
        yFemu[i] = fmu * sfr * run_params.qFesnia;
        yNimu[i] = fmu * sfr * run_params.qNisnia;
  }

  double yCr = integrate_arr(mu, yCrmu, count, max_mu/count, max_mu);
  double yFe = integrate_arr(mu, yFemu, count, max_mu/count, max_mu);
  double yNi = integrate_arr(mu, yNimu, count, max_mu/count, max_mu);
  for (i=0; i<count; i++){
        yCrphi[i] = yCr * compute_imf(mbin[i]);
        yFephi[i] = yFe * compute_imf(mbin[i]);
        yNiphi[i] = yNi * compute_imf(mbin[i]);
  //printf("SNIa done \n");
  }

  // AGB //
    if(run_params.AGBYields == 0)
    {
	i = j = 0;
	double Z_std[METALGRID] = {0.0, 1e-4, 4e-4, 4e-3, 8e-3, 0.02, 0.05};
	double sfrz = metallicity;
	if (sfrz < Z_std[0]){
		sfrz = Z_std[0];}
	else if (sfrz > Z_std[6]){
		sfrz = Z_std[6];}

	for (i=0; i<METALGRID; i++) {
		if(Z_std[i] <= sfrz && Z_std[i] > sfrz) {
			j = i;
		}
	}

	//double yCagb[run_params.countagb], yNagb[run_params.countagb], yOagb[run_params.countagb];
	//double m_agb[run_params.countagb];
	for(i=0; i<run_params.countagb; i++){
		if(run_params.magb[i] != 0) {
			m_agb[i] = run_params.magb[i];
			phi = compute_imf(run_params.magb[i]);
			taum = compute_taum(run_params.magb[i]);
			time = age[galaxies[p].SnapNum] - taum;
			if(time < run_params.lbtime[0]){
				sfr=0;
			}
			else {
				 sfr = interpolate_arr(age, sfh, SNAPLEN, time);
			}
			yCagb[i] = run_params.qCagb[i][j] * phi * sfr;
        		yNagb[i] = run_params.qNagb[i][j] * phi * sfr;
        		yOagb[i] = run_params.qOagb[i][j] * phi * sfr;
		}
		else {
        		yCagb[i] = 0;
        		yNagb[i] = 0;
        		yOagb[i] = 0;
		}
	}
    //printf("AGB done \n");
    }

  // SN II //
    if(run_params.SNIIYields == 0) //Woosley & Weaver 1995
    {
	i = j = 0;
	double Z_std[METALGRID] = {0.0, 1e-4, 4e-4, 4e-3, 8e-3, 0.02, 0.05};
	double sfrz = metallicity;
        if (sfrz < Z_std[0]){
                sfrz = Z_std[0];}
        else if (sfrz > Z_std[6]){
                sfrz = Z_std[6];}

        for (i=0; i<METALGRID; i++) {
                if(Z_std[i] <= sfrz && Z_std[i] > sfrz) {
                        j = i;
                }
        }

      	for(i=0; i<run_params.countsn; i++){
        	if(run_params.msn[i] != 0) {
         		m_sn[i] = run_params.msn[i];
         		phi = compute_imf(run_params.msn[i]);
          		taum = compute_taum(run_params.msn[i]);
          		time = age[galaxies[p].SnapNum] - taum;
          		if(time < run_params.lbtime[0]){
                		sfr=0;
          		}
         		 else {
                	sfr = interpolate_arr(age, sfh, SNAPLEN, time);
	  		}
          		yCsn[i] = run_params.qCsn[i][j] * phi * sfr;
          		yOsn[i] = run_params.qOsn[i][j] * phi * sfr;
          		yMgsn[i] = run_params.qMgsn[i][j] * phi * sfr;
          		ySisn[i] = run_params.qSisn[i][j] * phi * sfr;
          		ySsn[i] = run_params.qSsn[i][j] * phi * sfr;
          		yCasn[i] = run_params.qCasn[i][j] * phi * sfr;
          		yFesn[i] = run_params.qFesn[i][j] * phi * sfr;
        	}
        	else {
        	yCsn[i] = 0;
        	yOsn[i] = 0;
        	yMgsn[i] = 0;
        	ySisn[i] = 0;
        	ySsn[i] = 0;
        	yCasn[i] = 0;
        	yFesn[i] = 0;
        	}
    	}
     //printf("SNII done \n");
     }

    else if(run_params.SNIIYields == 1) //Nomoto et al. 2006
    {
        i = j = 0;
        double Z_std[METALGRID] = {0.0, 0.001, 0004, 0.02};
        double sfrz = metallicity;
        if (sfrz < Z_std[0]){
                sfrz = Z_std[0];}
        else if (sfrz > Z_std[3]){
                sfrz = Z_std[3];}

        for (i=0; i<METALGRID; i++) {
                if(Z_std[i] <= sfrz && Z_std[i] > sfrz) {
                        j = i;
                }
        }

        for(i=0; i<run_params.countsn; i++){
                if(run_params.msn[i] != 0) {
                        m_sn[i] = run_params.msn[i];
                        phi = compute_imf(run_params.msn[i]);
                        taum = compute_taum(run_params.msn[i]);
                        time = age[galaxies[p].SnapNum] - taum;
                        if(time < run_params.lbtime[0]){
                                sfr=0;
                        }
                         else {
                        sfr = interpolate_arr(age, sfh, SNAPLEN, time);
                        }
                        yCsn[i] = run_params.qCsn[i][j] * phi * sfr;
                        yOsn[i] = run_params.qOsn[i][j] * phi * sfr;
                        yMgsn[i] = run_params.qMgsn[i][j] * phi * sfr;
                        ySisn[i] = run_params.qSisn[i][j] * phi * sfr;
                        ySsn[i] = run_params.qSsn[i][j] * phi * sfr;
                        yCasn[i] = run_params.qCasn[i][j] * phi * sfr;
                        yFesn[i] = run_params.qFesn[i][j] * phi * sfr;
                }
                else {
                yCsn[i] = 0;
                yOsn[i] = 0;
                yMgsn[i] = 0;
                ySisn[i] = 0;
                ySsn[i] = 0;
                yCasn[i] = 0;
                yFesn[i] = 0;
                }
        }
     //printf("SNII done \n");
     }

    //printf("start integration");
    yCr_snia = A * integrate_arr(mbin, yCrphi, count, mbin[0], up_binary);
    yNi_snia = A * integrate_arr(mbin, yNiphi, count, mbin[0], up_binary);
    yFe_snia = A * integrate_arr(mbin, yFephi, count, mbin[0], up_binary);
    yC_agb = (1-A) * integrate_arr(m_agb, yCagb, run_params.countagb, low_binary, up_agb) + integrate_arr(m_agb, yCagb, run_params.countagb, low_agb, low_binary);
    yN_agb = (1-A) * integrate_arr(m_agb, yNagb, run_params.countagb, low_binary, up_agb) + integrate_arr(m_agb, yNagb, run_params.countagb, low_agb, low_binary);
    yO_agb = (1-A) * integrate_arr(m_agb, yOagb, run_params.countagb, low_binary, up_agb) + integrate_arr(m_agb, yOagb, run_params.countagb, low_agb, low_binary);
    yC_sn =  (1-A) * integrate_arr(m_sn, yCsn, run_params.countsn, low_sn, up_binary) + integrate_arr(m_sn, yCsn, run_params.countsn, up_binary, up_sn);
    yO_sn =  (1-A) * integrate_arr(m_sn, yOsn, run_params.countsn, low_sn, up_binary) + integrate_arr(m_sn, yOsn, run_params.countsn, up_binary, up_sn);
    yMg_sn =  (1-A) * integrate_arr(m_sn, yMgsn, run_params.countsn, low_sn, up_binary) + integrate_arr(m_sn, yMgsn, run_params.countsn, up_binary, up_sn);
    ySi_sn =  (1-A) * integrate_arr(m_sn, ySisn, run_params.countsn, low_sn, up_binary) + integrate_arr(m_sn, ySisn, run_params.countsn, up_binary, up_sn);
    yS_sn =  (1-A) * integrate_arr(m_sn, ySsn, run_params.countsn, low_sn, up_binary) + integrate_arr(m_sn, ySsn, run_params.countsn, up_binary, up_sn);
    yCa_sn =  (1-A) * integrate_arr(m_sn, yCasn, run_params.countsn, low_sn, up_binary) + integrate_arr(m_sn, yCasn, run_params.countsn, up_binary, up_sn);
    yFe_sn =  (1-A) * integrate_arr(m_sn, yFesn, run_params.countsn, low_sn, up_binary) + integrate_arr(m_sn, yFesn, run_params.countsn, up_binary, up_sn);
    //printf("end integration");

    double yield_agb = yC_agb + yN_agb + yO_agb;
    double yield_snia = yCr_snia + yNi_snia + yFe_snia;
    double yield_sn = yC_sn + yO_sn + yMg_sn + ySi_sn + yS_sn + yCa_sn + yFe_sn;
    yield = (yield_agb + yield_snia + yield_sn);

//overall metallicity
  if(galaxies[p].ColdGas > 5.0e-2) {
        galaxies[p].MetalsColdGas += yield * dt;
  }
  else {
        galaxies[centralgal].MetalsHotGas += yield * dt;
  }
  XPRINT(galaxies[p].MetalsColdGas <= galaxies[p].ColdGas, "metallicity = %.3f, cold gas mass = %.3e, galaxy ID = %i \n", galaxies[p].MetalsColdGas / galaxies[p].ColdGas, galaxies[p].ColdGas, galaxies[p].GalaxyNr);

//mass of each element formed in each production channel, should satisfy same condition with overall metallicity
  if(galaxies[p].ColdGas > 1.0e-2) {
        Cr_snia = yCr_snia * dt;
        Fe_snia = yFe_snia * dt;
        Ni_snia = yNi_snia * dt;
        C_agb = yC_agb * dt;
        N_agb = yN_agb * dt;
        O_agb = yO_agb * dt;
        C_sn = yC_sn * dt;
        O_sn = yO_sn * dt;
        Mg_sn = yMg_sn * dt;
        Si_sn = ySi_sn * dt;
        S_sn = yS_sn * dt;
        Ca_sn = yCa_sn * dt;
        Fe_sn = yFe_sn * dt;

        double dustdot = 0;

        //produce dust
        //double dust_agb, dust_sn, dust_snia;
        double delta_agb = 0.2;  //should be free parameter!! (value based on Popping et al. 2017)
        double delta_sn = 0.2;  //should be free parameter!! (value based on Popping et al. 2017)
        double delta_snia = 0.15;  //should be free parameter!! (value based on Popping et al. 2017)

        assert(dt > 0 && "dt must be greater than 0");

        /* eq 4 and eq 5 Popping et al. 2017 */
        //from AGB
        if (C_agb/O_agb > 1) {
           dustdot += delta_agb * (C_agb - 0.75*O_agb) / dt;
        }
        else {
           dustdot += delta_agb * (C_agb + N_agb + O_agb) / dt;
        }

        /* eq 6 Popping et al. 2017 */
        //from SN
        dustdot += delta_sn * C_sn / dt;
        dustdot += delta_sn * O_sn / dt;
        dustdot += 16 * delta_sn * (Mg_sn/24 + Si_sn/28 + S_sn/32 + Ca_sn/40 + Fe_sn/56) / dt;

        /* eq 6 Popping et al. 2017 */
        //from SNIa
        //printf("dustdot before = %f \n", dustdot);
	//dustdot += 16 * delta_snia * (Fe_snia/56) / dt;
        //dustdot += delta_snia * (Cr_snia + Ni_snia) / dt;

	//printf("dustdot after = %f \n", dustdot);
      	galaxies[p].ColdDust += dustdot * dt;
	galaxies[p].MetalsColdGas -= dustdot * dt;
        galaxies[p].dustdotform[step] += dustdot;
        XPRINT(dustdot * dt >= 0, "dust mass = %.3e, delta dust = %.3e, galaxy id = %i \n", galaxies[p].ColdDust, dustdot*dt, galaxies[p].GalaxyNr);

   }
}

void accrete_dust(const double metallicity, const double dt, const int p, const int step, struct GALAXY *galaxies) {
//dust accretion in ISM : eq 20 Asano13
  double dustdot = 0;
  double tacc_zero = 20 * SEC_PER_MEGAYEAR / run_params.UnitTime_in_s; //should be free parameter! yr

  if (galaxies[p].MetalsColdGas > galaxies[p].ColdDust && metallicity > 0  ) {
    double tacc = tacc_zero * 0.02 / metallicity;
    dustdot += (1 - galaxies[p].ColdDust/galaxies[p].MetalsColdGas) * (galaxies[p].f_H2 * galaxies[p].ColdDust / tacc);

    galaxies[p].dustdotgrowth[step] += dustdot;
    galaxies[p].ColdDust += dustdot * dt;
    galaxies[p].MetalsColdGas -= dustdot * dt;
    XPRINT(galaxies[p].ColdDust >= 0, "dust mass = %.3e, delta dust = %.3e, galaxy id = %i \n", galaxies[p].ColdDust, dustdot*dt, galaxies[p].GalaxyNr);
  }

}

void destruct_dust(const double metallicity, const double stars, const double dt, const int p, const int step,  struct GALAXY *galaxies) {
//dust destruction : Asano et al. 13
  int i;
  double sfr;
  double dustdot = 0;
  double age[SNAPLEN], sfh[SNAPLEN];
  double m_low = 8; //lower limit of stellar mass that end up as SN
  double up_binary = 16; //upper limit for binary stars, Msun
  double m_up = 40; //upper limit of stellar mass that end up as SN
  double t_low = compute_taum(m_up); //the shortest stellar lifetime that end up as SN
  double eta = 0.1; //should be free parameter, we used value adopted by Asano et al. 13
  double m_swept = 1535 * pow((metallicity/0.02 + 0.039), -0.289) * run_params.Hubble_h / 1.e10; //eq. 14 Asano et al. 13
  int count = 20;
  double mass[count], phi[count], mphi[count];
  double A = run_params.BinaryFraction;
//  double m_swept = 600 * run_params.Hubble_h / 1.e10;

  for (i=0; i<SNAPLEN; i++)
  {
        age[i] = run_params.lbtime[i]; //change the variable name to represent the age of universe instead of lbtime
        sfh[i] = galaxies[p].Sfr[i];
  }

//Dwek & Cherchneff 2011
  if (age[galaxies[p].SnapNum] > t_low) {
    for(i=0; i<count; i++) {
        mass[i] = 1 + ((m_up - 1) / (count-i));
        phi[i] = compute_imf(mass[i]);
        mphi[i] = mass[i]*phi[i];
    }
/*
    for (i=0; i<count; i++){
	mass[i] = m_low + ((m_up - m_low) / (count-i));
	phi[i] = compute_imf(mass[i]);
        double taum = compute_taum(mass[i]);
        double time = age[galaxies[p].SnapNum] - taum;
        if(time < run_params.lbtime[0]){
                sfr=0;
        }
        else {
                sfr = interpolate_arr(age, sfh, SNAPLEN, time);
        }
	mphi[i] = sfr*phi[i];
    }
*/

    double mstar = integrate_arr(mass, mphi, count, mass[0], m_up) / integrate_arr(mass, phi, count, m_low, m_up);
    double mstar_ia = integrate_arr(mass, mphi, count, mass[0], m_up) / integrate_arr(mass, phi, count, m_low, up_binary);
    double mstar_ii = integrate_arr(mass, mphi, count, mass[0], m_up) / integrate_arr(mass, phi, count, up_binary, m_up);
    double Rsn_ia = A*stars / dt / mstar_ia;
    double Rsn_ii = (1-A) * stars / dt / mstar_ia + stars/dt/mstar_ii;
    //printf("%f \n", Rsn_ii/Rsn_ia);
    double Rsn = stars / dt / mstar;
    assert(m_swept > 0 && "mass of ISM swept by SN must be greater than 0");

//    if (Rsn > 0 && galaxies[p].ColdGas > 0 && galaxies[p].f_HI >0) {
    if (Rsn > 0 && galaxies[p].ColdGas > 0) {
       Rsn *=  run_params.UnitMass_in_g / run_params.UnitTime_in_s * SEC_PER_YEAR / SOLAR_MASS; //convert to Msun/yr
//       double tsn = galaxies[p].ColdGas * galaxies[p].f_HI/ (eta * m_swept * Rsn); //eq.12 Asano et al 13 [yr]
       double tsn = galaxies[p].ColdGas / (eta * m_swept * Rsn); //eq.12 Asano et al 13 [yr]
       tsn /= run_params.UnitTime_in_s / SEC_PER_YEAR; //back to internal unit
//       printf("coldgas = %.3e, dust = %.3e, m_swept = %.3e, Rsn = %.3e,  tsn = %.3e, dt = %.3e Myr \n", galaxies[p].ColdGas, galaxies[p].ColdDust,  m_swept, Rsn, tsn * run_params.UnitTime_in_s / SEC_PER_MEGAYEAR, dt * run_params.UnitTime_in_s / SEC_PER_MEGAYEAR);
//       assert(tsn > 0 && "tsn must be greater than 0");
       dustdot += galaxies[p].ColdDust / tsn;
    }

    if (galaxies[p].ColdDust - dustdot * dt > 0) {
  	  galaxies[p].dustdotdestruct[step] += dustdot;
  	  galaxies[p].ColdDust -= dustdot * dt;
  	  galaxies[p].MetalsColdGas += dustdot * dt;
    }
    else {
	galaxies[p].dustdotdestruct[step] += dustdot;
	galaxies[p].MetalsColdGas += galaxies[p].ColdDust;
	galaxies[p].ColdDust = 0;
    }

  }
}

void dust_thermal_sputtering(const int gal, const double dt, struct GALAXY *galaxies) //section 3.5 Popping 2017
{
   double adot;
   double a = 1e-5 / run_params.UnitLength_in_cm; //a = 0.1 micrometer
   double mdot = 0.0;

// thermal sputtering in hot dust
   if(galaxies[gal].HotDust > 0.0 && galaxies[gal].Vvir > 0.0) {
        const double temp0 = 2e6; //in Kelvin
        const double temp = 35.9 * galaxies[gal].Vvir * galaxies[gal].Vvir; //in Kelvin
        const double gamma = 2.5;

        double x = -3.2e-18 / PROTONMASS / (pow(temp0/temp, gamma) + 1); //in cm^4 / g /s
        x /= run_params.UnitLength_in_cm / run_params.UnitDensity_in_cgs / run_params.UnitTime_in_s; //in internal unit
        double rho = galaxies[gal].HotGas / (4 / 3 * M_PI * galaxies[gal].Rvir * galaxies[gal].Rvir * galaxies[gal].Rvir);
        adot = rho / x;
        a -= adot * dt;

        assert(a > 0);
        a /= 1e-5 / run_params.UnitLength_in_cm; //in 0.1 micrometer unit
        rho *= run_params.UnitDensity_in_cgs;
        const double tau0 = 170 * SEC_PER_MEGAYEAR / run_params.UnitTime_in_s; // 0.17Gyr to internal unit
        const double tau = tau0 * (a/rho) * (pow(temp0/temp, gamma) + 1);
        mdot = galaxies[gal].HotDust / tau / 3;

        galaxies[gal].HotDust -= mdot * dt;
        galaxies[gal].MetalsHotGas += mdot * dt;
   }

// thermal sputtering in ejected dust
   if(galaxies[gal].EjectedDust > 0.0 && galaxies[gal].Vvir > 0.0) {
        const double temp0 = 2e6; //in Kelvin
        const double temp = 35.9 * galaxies[gal].Vvir * galaxies[gal].Vvir; //in Kelvin
        const double gamma = 2.5;

        double x = -3.2e-18 / PROTONMASS / (pow(temp0/temp, gamma) + 1); //in cm^4 / g /s
        x /= run_params.UnitLength_in_cm / run_params.UnitDensity_in_cgs / run_params.UnitTime_in_s; //in internal unit
        double rho = galaxies[gal].EjectedMass / (4 / 3 * M_PI * galaxies[gal].Rvir * galaxies[gal].Rvir * galaxies[gal].Rvir);
        adot = rho / x;
        a -= adot * dt;

        assert(a > 0);
        a /= 1e-5 / run_params.UnitLength_in_cm; //in 0.1 micrometer unit
        rho *= run_params.UnitDensity_in_cgs;
        const double tau0 = 170 * SEC_PER_MEGAYEAR / run_params.UnitTime_in_s; // 0.17Gyr to internal unit
        const double tau = tau0 * (a/rho) * (pow(temp0/temp, gamma) + 1);
        mdot = galaxies[gal].EjectedDust / tau / 3;

        galaxies[gal].EjectedDust -= mdot * dt;
        galaxies[gal].MetalsEjectedMass += mdot * dt;
   }

}

double integrate_arr(const double arr1[MAX_STRING_LEN], const double arr2[MAX_STRING_LEN], const int npts, const double lower_limit, const double upper_limit)
{
        double Q;
        gsl_interp_accel *acc;
        gsl_spline *spl;

        acc = gsl_interp_accel_alloc ();
        spl = gsl_spline_alloc (gsl_interp_linear, npts);

        gsl_spline_init (spl, arr1, arr2, npts);
        Q = gsl_spline_eval_integ (spl, lower_limit, upper_limit, acc);

        gsl_spline_free (spl);
        gsl_interp_accel_free (acc);

        return Q;
}

double interpolate_arr(const double arr1[MAX_STRING_LEN], const double arr2[MAX_STRING_LEN], const int npts, const int xi)
{
        double Q;
        gsl_interp_accel *acc;
        gsl_spline *spl;

        acc = gsl_interp_accel_alloc ();
        spl = gsl_spline_alloc (gsl_interp_linear, npts);

        gsl_spline_init (spl, arr1, arr2, npts);
        Q = gsl_spline_eval (spl, xi, acc);

        gsl_spline_free (spl);
        gsl_interp_accel_free (acc);
        return Q;
}

double compute_imf (const double m)
{ //eq 11 Arrigoni et al. 2010
        double mass = m;
        double A = 0.9098, B = 0.2539, x = 1.3, sigma=0.69;
        double mc = 0.079; //Msun
        double phi;

        if (m < 1){
                phi = A * exp( -(log10(mass) - log10(mc)) / (2*sigma*sigma));
                }
        else {
                phi = B * pow(mass, -x);
                }

        return phi;
}

double compute_taum (const double m)
{ //eq 3 Raiteri et al. 1996
        double mass = m;
        double Z = 0.02;
        double a0 = 10.13 + 0.07547*log10(Z) - 0.008084*log10(Z)*log10(Z);
        double a1 = -4.424 - 0.7939*log10(Z) - 0.1187*log10(Z)*log10(Z);
        double a2 = 1.262 + 0.3385*log10(Z) + 0.05417*log10(Z)*log10(Z);

        double logt = a0 + a1*log10(mass) + a2*log10(mass)*log10(mass);
        double t = pow(10, logt) / 1e6  ; //in Myr/h

        /* convert into Myrs/h */
        return t;
}

max_stars_repo_name: dptriani/sage

max_issues_repo_name: dptriani/sage

max_forks_repo_name: dptriani/sage


text:
/*
 * Copyright (c) 2018-2021 Aleksas Mazeliauskas, Stefan Floerchinger,
 *                    Eduardo Grossi, and Derek Teaney
 * All rights reserved.
 *
 * FastReso is distributed under MIT license;
 * see the LICENSE file that should be present in the root
 * of the source distribution, or alternately available at:
 * https://github.com/amazeliauskas/FastReso/
 */
#ifndef FASTRESO_TFastReso_formulas_h
#define FASTRESO_TFastReso_formulas_h
#include <gsl/gsl_integration.h>
#include "TParticle.h"
#include <vector>
// List of parameters for GSL integration procedure.
#include "qag_params.h"
#include "grid_params.h"

#include <gsl/gsl_errno.h>
//using namespace std;
//! Eabc -- energy of b in restframe of a in a->b+c decay
double get_Eabc(double ma,double mb,double mc) { return (ma*ma+mb*mb-mc*mc)/(2*ma);}
//! particle momenta
double get_p(double E, double m) { return sqrt(E*E-m*m);}
//! pabc -- momentum of b or c in restframe of a in a->b+c decay
double get_pabc(double ma,double mb,double mc) { return get_p(get_Eabc(ma,mb,mc),mb);}

//! returns thermal Bose-Einstein, Fermi-Dirac or Boltzmann distribution
double get_thermal_F(double Ebar, double T, double QMu, EParticleType type);
//! returns derivative of thermal Bose-Einstein, Fermi-Dirac or Boltzmann distribution
double get_thermal_dF(double Ebar, double T, double QMu, EParticleType type);

//! Initialize f_j on the freeze-out surface
//! Here you can add additional fj's
double get_initial_Fj(EFjIndex j, double Ebar, double T, double QMu, EParticleType type, double m, double Cs2)
{
  double pbar = sqrt(Ebar*Ebar-m*m);
  switch(j){
    case EFjIndex::kFeq1: //feq 1
      return pbar*get_thermal_F(Ebar,T, QMu, type);
      break;
    case EFjIndex::kFeq2: //feq 2
      return pbar*get_thermal_F(Ebar,T, QMu, type);
      break;
    case EFjIndex::kFshear1: //fshear 1
      return pbar*pbar*pbar*get_thermal_dF(Ebar,T, QMu, type);
      break;
    case EFjIndex::kFshear2: //fshear 2
      return pbar*pbar*pbar*get_thermal_dF(Ebar,T, QMu, type);
      break;
    case EFjIndex::kFshear3: //fshear 3
      return pbar*pbar*pbar*get_thermal_dF(Ebar,T, QMu, type);
      break;
    case EFjIndex::kFbulk1: //fbulk 1
      return pbar*get_thermal_dF(Ebar,T, QMu, type)*(1./3.*m*m/T/Ebar-Ebar/T*(1./3.-Cs2));
      break;
    case EFjIndex::kFbulk2: //fbulk 2
      return pbar*get_thermal_dF(Ebar,T, QMu, type)*(1./3.*m*m/T/Ebar-Ebar/T*(1./3.-Cs2));
      break;
    case EFjIndex::kFtemp1: //ftemperature. 1
      return pbar*get_thermal_dF(Ebar,T, QMu, type)*Ebar/T;
      break;
    case EFjIndex::kFtemp2: //ftemperature 2
      return pbar*get_thermal_dF(Ebar,T, QMu, type)*Ebar/T;
      break;
    case EFjIndex::kFvel1: //fvelocity  1
      return pbar*pbar*get_thermal_dF(Ebar,T, QMu, type);
      break;
    case EFjIndex::kFvel2: //fvelocity 2
      return pbar*pbar*get_thermal_dF(Ebar,T, QMu, type);
      break;
    case EFjIndex::kFvel3: //fvelocity 3
      return pbar*pbar*get_thermal_dF(Ebar,T, QMu, type);
      break;
    default:
      std::cerr  << "\033[1mTFastReso.cpp::get_initial_Fj\033[0m : \033[1;31merror\033[0m :  wrong index "  << (int) j << std::endl;
      exit(EXIT_FAILURE);
  }
}
//! Transformation rule factor A for Fj in a 2-body decay
//! Here you can add additional rules for fj's
double get_factor_Aj(EFjIndex j, double Qw, double Ew, double Ebar, double pbar, double Ma, double pw)
{

//  double A= (Ma*Eabc/Mb/Mb-w*Ma*pabc/Mb/Mb*Ebar/pbar);
//  double Ew= (Ma*Eabc*Ebar/Mb/Mb-w*Ma*pabc*pbar/Mb/Mb);

  switch( j){
    case EFjIndex::kFeq1: // feq 1
      return Qw/pw;
      break;
    case EFjIndex::kFeq2: // feq 2
      return pbar/pw*Ew/Ebar;
      break;
    case EFjIndex::kFshear1: // fshear 1
      return Qw/pw*(2.5*Qw/pw*Qw/pw-1.5);
      break;
    case EFjIndex::kFshear2: // fshear 2
      return Qw/pw;
      break;
    case EFjIndex::kFshear3: // fshear 3
      return (1.5*Qw/pw*Qw/pw-0.5)*Ew/Ebar*pbar/pw;
      break;
    case EFjIndex::kFbulk1: // fbulk 1
      return Qw/pw;
      break;
    case EFjIndex::kFbulk2: // fbulk 2
      return pbar/pw*Ew/Ebar;
      break;
    case EFjIndex::kFtemp1: // ftemperature 1
      return Qw/pw;
      break;
    case EFjIndex::kFtemp2: // ftemperature 2
      return pbar/pw*Ew/Ebar;
      break;
    case EFjIndex::kFvel1: // fvelocity 1
      return 1.5*Qw/pw*Qw/pw-0.5;
      break;
    case EFjIndex::kFvel2: // fvelocity 2
      return 1;
      break;
    case EFjIndex::kFvel3: // fvelocity 3
      return pbar/pw*Qw/pw*Ew/Ebar;
      break;
    default:
      std::cerr  << "\033[1mTFastReso.cpp::get_factor_Aj\033[0m : \033[1;31merror\033[0m :  wrong index "  << (int) j << std::endl;
      exit(EXIT_FAILURE);
  }
}

//! Parameters for integration of thermal distribution
struct thermal_params{
  double m;
  double Tfo;
  double QMu;
  EParticleType type;
};
//! Thermal distribution function for integration over vbar=[0,1]
double get_F_p_dvbar(double vbar, void * p) {
  double &m           = ((struct thermal_params *) p)->m;
  double &Tfo           = ((struct thermal_params *) p)->Tfo;
  double &QMu           = ((struct thermal_params *) p)->QMu;
  EParticleType &type           = ((struct thermal_params *) p)->type;
  double Ebar=m/sqrt(1-vbar*vbar);
  double pbar=vbar*Ebar;
  double F_p = get_thermal_F(Ebar,Tfo,QMu, type);
  return F_p*pbar*pbar*4*M_PI*m/pow(2*M_PI*sqrt(1-vbar*vbar),3);
}
//! Parameters for integration of two body decay function
struct twobody_params{
  double Ebar;
  double Ma;
  double Mb;
  double Mc;
  TParticle * Parent;
  EFjIndex index;
};
//! Integrand of two body decay formula for f_j components
double get_F_pu_dw(double w, void * p) {
  double &Ebar           = ((struct twobody_params *) p)->Ebar;
  double &Ma             = ((struct twobody_params *) p)->Ma;
  double &Mb             = ((struct twobody_params *) p)->Mb;
  double &Mc             = ((struct twobody_params *) p)->Mc;
  TParticle *Parent      = ((struct twobody_params *) p)->Parent;
  EFjIndex &index             = ((struct twobody_params *) p)->index;

  double Eabc = get_Eabc(Ma,Mb,Mc);
  double pabc = get_p(Eabc, Mb);
  //double pbar = GSL_MAX(get_p(Ebar, Mb),DBL_EPSILON);
  double pbar = get_p(Ebar, Mb);// ,DBL_EPSILON);
  double Ebar_old = Ebar*Eabc*Ma/Mb/Mb-w*pbar*pabc*Ma/Mb/Mb;
  //double Ebar_old = pbar*pabc*Ma/Mb/Mb*(1-w);
  //double Ebar_old = pbar*pabc/Ma*u;
  double F_old=Parent->get_Fj( (int) index, Ebar_old);
  double Qw= (Ma*Eabc*pbar/Mb/Mb-w*Ma*pabc/Mb/Mb*Ebar);
  //double Qw = pabc/Ma*u;
//  double Ew= (Ma*Eabc*Ebar/Mb/Mb-w*Ma*pabc*pbar/Mb/Mb);
  double A = get_factor_Aj(index, Qw, Ebar_old, Ebar, pbar, Ma,sqrt(Qw*Qw+(1-w*w)*Ma*Ma/Mb/Mb*pabc*pabc));
  return F_old*A;
}
double get_F_pu_du(double u, void * p) {
  double &Ebar           = ((struct twobody_params *) p)->Ebar;
  double &Ma             = ((struct twobody_params *) p)->Ma;
  double &Mb             = ((struct twobody_params *) p)->Mb;
  double &Mc             = ((struct twobody_params *) p)->Mc;
  TParticle *Parent      = ((struct twobody_params *) p)->Parent;
  EFjIndex &index             = ((struct twobody_params *) p)->index;

  double Eabc = get_Eabc(Ma,Mb,Mc);
  double pabc = get_p(Eabc, Mb);
  //double pbar = GSL_MAX(get_p(Ebar, Mb),DBL_EPSILON);
  double pbar = get_p(Ebar, Mb);// ,DBL_EPSILON);
  //double Ebar_old = Ebar*Eabc*Ma/Mb/Mb-w*pbar*pabc*Ma/Mb/Mb;
  //double Ebar_old = pbar*pabc*Ma/Mb/Mb*(1-w);
  if (pbar==0) { return 0;};
  double Ebar_old = Ma/2*(pbar/pabc+pabc/pbar)+u*Ma;
  double F_old=Parent->get_Fj((int) index, Ebar_old);
  //double Qw= (Ma*Eabc/Mb/Mb-w*Ma*pabc/Mb/Mb*Ebar/pbar);
  double Qw = Ma/2/pbar*(pbar*pbar/pabc-pabc)+Ma*u;
//  double Ew= (Ma*Eabc*Ebar/Mb/Mb-w*Ma*pabc*pbar/Mb/Mb);
  double A = get_factor_Aj(index, Qw, Ebar_old, Ebar, pbar, Ma,sqrt(Qw*Qw+2*u*pabc/pbar*Ma*Ma));
  return F_old*A/(pbar*pabc/Ma/Ma);
}

//! Parameters for three body decay mass integral
struct threebody_params{
  gsl_function  *twobody_function;
  double Mc;
  double Md;
  gsl_integration_workspace *twobody_workspace;
};

//! Integrand for three body decay normalization
double get_B_dm(double mct, void * p) {
  gsl_function *twobody_function = ((struct threebody_params *) p)->twobody_function;
  double &Ma = ((struct twobody_params *) twobody_function->params)->Ma;
  double &Mb = ((struct twobody_params *) twobody_function->params)->Mb;
  double &Mc             = ((struct threebody_params *) p)->Mc;
  double &Md             = ((struct threebody_params *) p)->Md;
  return get_pabc(Ma,Mb,mct)*get_pabc(mct, Mc, Md);
}
//! Integrand for three body decay function
double get_F_pu_dm(double mct, void * p) {
  gsl_function *twobody_function = ((struct threebody_params *) p)->twobody_function;
  ((struct twobody_params *) twobody_function->params)->Mc = mct;
  gsl_integration_workspace *twobody_workspace = ((struct threebody_params *) p)->twobody_workspace;
  double &Ma = ((struct twobody_params *) twobody_function->params)->Ma;
  double &Mb = ((struct twobody_params *) twobody_function->params)->Mb;
  double &Mc             = ((struct threebody_params *) p)->Mc;
  double &Md             = ((struct threebody_params *) p)->Md;
  double result,error;

  using namespace qag_params;
  int status = gsl_integration_qag (twobody_function, -1, 1, fEpsAbs,fEpsRel, fLimit,fKey ,twobody_workspace, &result, &error);
  // Check status of integration and reduce accuracy if failing
  if (status) {
    status = gsl_integration_qag (twobody_function, -1, 1, fEpsAbs,1e-4, fLimit,GSL_INTEG_GAUSS15,twobody_workspace, &result, &error);
    if (status) {
      status = gsl_integration_qag (twobody_function, -1, 1, 1e-4,1e-4, fLimit,GSL_INTEG_GAUSS15,twobody_workspace, &result, &error);
      if (status) {      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31merror\033[0m : integration failure! status = " <<gsl_strerror ( status )<< ". " << result << "+_" << error << std::endl;
        exit(EXIT_FAILURE);
      }else {
        std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative and absolute error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
      }
    } else {
      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
    }
  }
  return result*get_pabc(Ma,Mb,mct)*get_pabc(mct, Mc,Md);
}
double get_F_pu_dm_u(double mct, void * p) {
  gsl_function *twobody_function = ((struct threebody_params *) p)->twobody_function;
  ((struct twobody_params *) twobody_function->params)->Mc = mct;
  gsl_integration_workspace *twobody_workspace = ((struct threebody_params *) p)->twobody_workspace;
  double &Ma = ((struct twobody_params *) twobody_function->params)->Ma;
  double &Mb = ((struct twobody_params *) twobody_function->params)->Mb;
  double &Mc             = ((struct threebody_params *) p)->Mc;
  double &Md             = ((struct threebody_params *) p)->Md;
  double result,error;

  using namespace qag_params;
  int status = gsl_integration_qagiu (twobody_function, 0.0, fEpsAbs,fEpsRel, fLimit,twobody_workspace, &result, &error);
  // Check status of integration and reduce accuracy if failing
  if (status) {
  status = gsl_integration_qagiu (twobody_function, 0.0, fEpsAbs,1e-4, fLimit,twobody_workspace, &result, &error);
    if (status) {
  status = gsl_integration_qagiu (twobody_function, 0.0, 1e-4,1e-4, fLimit,twobody_workspace, &result, &error);
      if (status) {      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31merror\033[0m : integration failure! status = " << gsl_strerror (status )<< ". " << result << "+_" << error << std::endl;
        exit(EXIT_FAILURE);
      }else {
        std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative and absolute error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
      }
    } else {
      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
    }
  }

  return result*get_pabc(Ma,Mb,mct)*get_pabc(mct, Mc,Md);
}

void do_2bodydecay(TParticle * Parent, TParticle * Child1, TParticle * Child2, double BranchingRatio);
void do_3bodydecay(TParticle * Parent, TParticle * Child1, TParticle * Child2, TParticle * Child3, double BranchingRatio);

//! Initialize thermal distribution
double get_thermal_F(double Ebar, double T, double QMu, EParticleType type)
{
  switch(type){
    case EParticleType::kBoson:
      return 1.0/(exp((Ebar-QMu)/T)-1);
      break;
    case EParticleType::kFermion:
      return 1.0/(exp((Ebar-QMu)/T)+1);
      break;
    case EParticleType::kBoltzman:
      return exp(-(Ebar-QMu)/T);
      break;
    default:
      std::cerr << "\033[1mTFastReso.cpp::get_thermal_F\033[0m : \033[1;31merror\033[0m : unrecognined particle type!" <<std::endl;
      exit(EXIT_FAILURE);}
}
//! Initialize  derivative of t hermal distribution (With a minus sign)
double get_thermal_dF(double Ebar, double T, double QMu, EParticleType type)
{
  switch(type){
    case EParticleType::kBoson:
      return 1.0/(exp((Ebar-QMu)/T)-1)*(1+1.0/(exp((Ebar-QMu)/T)-1));
      break;
    case EParticleType::kFermion:
      return 1.0/(exp((Ebar-QMu)/T)+1)*(1-1.0/(exp((Ebar-QMu)/T)+1));
      break;
    case EParticleType::kBoltzman:
      return exp(-(Ebar-QMu)/T);
      break;
    default:
      std::cerr << "\033[1mTFastReso.cpp::get_thermal_F_pu\033[0m : \033[1;31merror\033[0m : unrecognined particle type!" <<std::endl;
      exit(EXIT_FAILURE);}
}

//! Perform two body decay integral: Parent->Child1 + Child2
void do_2bodydecay(TParticle * Parent, TParticle * Child1, TParticle * Child2, double BranchingRatio){
  using namespace qag_params;

  gsl_integration_workspace * fWorkspace2body ;
  fWorkspace2body = gsl_integration_workspace_alloc (fLimit);
  Parent->lock();
  double Ma = Parent->getM();
  double Mb = Child1->getM();
  double Mc = Child2->getM();
  double nua = Parent->getNu();
  double nub = Child1->getNu();
  double Ebar=0.0;
  Child1->clean_buffer();
  if (Ma < Mb+Mc) {
    std::cerr  << "\033[1mTFastReso.cpp::do_2bodydecays\033[0m : \033[1;31merror\033[0m : father particle lighter than childern! " << Ma << " < " << Mb+Mc << std::endl;
    exit(EXIT_FAILURE);
  }
  gsl_function func2body;
      if (Mb==0) {
  func2body.function = get_F_pu_du;
      } else {
  func2body.function = get_F_pu_dw;
      }
  twobody_params params = {Ebar, Ma, Mb, Mc, Parent, EFjIndex::kFeq1};
  func2body.params = &params;
  double error, result ;
  int sym=1 ;
  if (Child1==Child2 ){ sym=2; } else { sym=1; }

  double Nfather = Parent->getN();
  //! add fraction of the total yield to child
  Child1->addN(BranchingRatio*Nfather*sym);
  for (int i=0; i < Child1->getNpbar(); i++){
    double Ebar =Child1->getEbar(i);
    params.Ebar =Ebar;

    //for (int j=0; j < grid_params::fNf ; j++){
    for(auto j: Parent->fComponents) {
      params.index =j;
      double fac = 0;
      if (Mb==0) {
int status = gsl_integration_qagiu (&func2body, 0.0, fEpsAbs,fEpsRel, fLimit ,fWorkspace2body, &result, &error);

  // Check status of integration and reduce accuracy if failing
  if (status) {
status = gsl_integration_qagiu (&func2body, 0.0, fEpsAbs,1e-4, fLimit ,fWorkspace2body, &result, &error);
    if (status) {
status = gsl_integration_qagiu (&func2body, 0.0, 1e-4,1e-4, fLimit ,fWorkspace2body, &result, &error);
      if (status) {      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31merror\033[0m : integration failure! status = " << gsl_strerror (status )<< ". " << result << "+_" << error << std::endl;
        exit(EXIT_FAILURE);
      }else {
        std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative and absolute error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
      }
    } else {
      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
    }
  }

      fac = sym*BranchingRatio*nua/nub/2;
      }
      else {
int status = gsl_integration_qag (&func2body, -1, 1, fEpsAbs,fEpsRel, fLimit,fKey ,fWorkspace2body, &result, &error);

  // Check status of integration and reduce accuracy if failing
  if (status) {
status = gsl_integration_qag (&func2body, -1, 1, fEpsAbs,1e-4, fLimit, GSL_INTEG_GAUSS15,fWorkspace2body, &result, &error);
    if (status) {
status = gsl_integration_qag (&func2body, -1, 1, 1e-4,1e-4, fLimit, GSL_INTEG_GAUSS15 ,fWorkspace2body, &result, &error);
      if (status) {      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31merror\033[0m : integration failure! status = " << gsl_strerror (status) << ". " << result << "+_" << error << std::endl;
        exit(EXIT_FAILURE);
      }else {
        std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative and absolute error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
      }
    } else {
      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
    }
  }

      fac = sym*BranchingRatio*nua/nub*pow(Ma/Mb,2)/2;

      }
      Child1->addFj((int) j,i, fac*result);

    }
  }
  //if (Child1=="pi0139plu"  ) { getParticle(Child1)->print_buffer(Child1+"_"+Parent+"_"+Child2);}
  gsl_integration_workspace_free (fWorkspace2body);

}

//! Perform three body decay integral Parent->Child1 + Child2 + Child3
void do_3bodydecay(TParticle * Parent, TParticle * Child1, TParticle * Child2, TParticle * Child3, double BranchingRatio){
  using namespace qag_params;
  gsl_integration_workspace * fWorkspace2body ;
  gsl_integration_workspace * fWorkspace3body ;
  fWorkspace2body = gsl_integration_workspace_alloc (fLimit);
  fWorkspace3body = gsl_integration_workspace_alloc (fLimit);
  Parent->lock();
  double Ma = Parent->getM();
  double Mb = Child1->getM();
  double Mc = Child2->getM();
  double Md = Child3->getM();
  double nua = Parent->getNu();
  double nub = Child1->getNu();
  double Ebar=0.0;

  Child1->clean_buffer();
  if (Ma < Mb+Mc+Md) {
    std::cerr  << "\033[1mTFastReso.cpp::do_3bdoydecays\033[0m : \033[1;31merror\033[0m : father particle lighter than childern! " << Ma << " < " << Mb+Mc+Md << std::endl;
    exit(EXIT_FAILURE);
  }
  gsl_function func2body;
  gsl_function func3body;
      if (Mb==0) {
  func2body.function = get_F_pu_du;
      } else {
  func2body.function = get_F_pu_dw;
      }
  func3body.function = get_B_dm;
  twobody_params params2 = {Ebar, Ma, Mb, Mc+Md, Parent, EFjIndex::kFeq1};
  func2body.params = &params2;
  threebody_params params3 = {&func2body, Mc,Md, fWorkspace2body};
  func3body.params = &params3;
  double error, result ;
  int status =gsl_integration_qag (&func3body, Mc+Md, Ma-Mb, fEpsAbs,fEpsRel, fLimit,fKey ,fWorkspace3body, &result, &error);
  // Check status of integration and reduce accuracy if failing
  if (status) {
  status =gsl_integration_qag (&func3body, Mc+Md, Ma-Mb, fEpsAbs,1e-4, fLimit,GSL_INTEG_GAUSS15 ,fWorkspace3body, &result, &error);
    if (status) {
  status =gsl_integration_qag (&func3body, Mc+Md, Ma-Mb, 1e-4,1e-4, fLimit,GSL_INTEG_GAUSS15 ,fWorkspace3body, &result, &error);
      if (status) {      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31merror\033[0m : integration failure! status = " << gsl_strerror (status )<< ". " << result << "+_" << error << std::endl;
        exit(EXIT_FAILURE);
      }else {
        std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative and absolute error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
      }
    } else {
      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
    }
  }

  double B3body = result;
      if (Mb==0) {
  func3body.function = get_F_pu_dm_u;
      } else {
  func3body.function = get_F_pu_dm;
      }

  int sym = 1 ;
  if (Child1==Child2 and Child2==Child3 ){ sym=3; } else if (Child1==Child2 or Child1==Child3 ){ sym=2;} else { sym=1; }
  double Nfather = Parent->getN();
  // add fraction of yield to child
  Child1->addN(BranchingRatio*Nfather*sym);
  for (int i=0; i < Child1->getNpbar(); i++){
    double Ebar =Child1->getEbar(i);
    params2.Ebar =Ebar;
    //for (int j=0; j < grid_params::fNf ; j++){
    for(auto j: Parent->fComponents) {
      params2.index =j;
      int status = gsl_integration_qag (&func3body, Mc+Md, Ma-Mb, fEpsAbs,fEpsRel, fLimit,fKey ,fWorkspace3body, &result, &error);
  // Check status of integration and reduce accuracy if failing
  if (status) {
      status = gsl_integration_qag (&func3body, Mc+Md, Ma-Mb, fEpsAbs,1e-4, fLimit,GSL_INTEG_GAUSS15 ,fWorkspace3body, &result, &error);
    if (status) {
      status = gsl_integration_qag (&func3body, Mc+Md, Ma-Mb, 1e-4,1e-4, fLimit,GSL_INTEG_GAUSS15 ,fWorkspace3body, &result, &error);
      if (status) {      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31merror\033[0m : integration failure! status = " << gsl_strerror (status )<< ". " << result << "+_" << error << std::endl;
        exit(EXIT_FAILURE);
      }else {
        std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative and absolute error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
      }
    } else {
      std::cerr << "\033[1mTFastReso_formulas.h\033[0m : \033[1;31mwarning\033[0m : reduced relative error (1e-4) integration. Result = " << result << "+-" << error <<std::endl;
    }
  }

      double fac =0;
      if (Mb==0) {
        fac = sym*BranchingRatio*nua/nub/2/B3body;
      }
      else {
        fac = sym*BranchingRatio*nua/nub*pow(Ma/Mb,2)/2/B3body;
      }
      Child1->addFj((int) j,i, fac*result);

    }
  }

  //if (Child1=="pi0139plu" ) { getParticle(Child1)->print_buffer(Child1+"_"+Parent+"_"+Child2+"_"+Child3);}
  gsl_integration_workspace_free (fWorkspace2body);
  gsl_integration_workspace_free (fWorkspace3body);
}

#endif

max_stars_repo_name: amazeliauskas/FastReso

max_issues_repo_name: amazeliauskas/FastReso

max_forks_repo_name: amazeliauskas/FastReso


text:
#pragma once

#include "common.h"

#include <gsl/gsl_rng.h>

#define ALT_CNT 4

enum categorical_flags {
    TEST_TYPE_CODOMINANT = 1,
    TEST_TYPE_RECESSIVE = 2,
    TEST_TYPE_DOMINANT = 4,
    TEST_TYPE_ALLELIC = 8
};

struct categorical_supp {
    uint8_t *phen_bits;
    size_t *filter, *table, *phen_mar, *outer;
};

struct categorical_res {
    double nlpv[ALT_CNT], qas[ALT_CNT];
};

struct maver_adj_supp {
    uint8_t *phen_bits;
    size_t *filter, *table, *phen_mar, *phen_perm, *outer;
    struct categorical_snp_data *snp_data;
};

struct maver_adj_res {
    double nlpv[ALT_CNT];
    size_t rpl[ALT_CNT];
};

double stat_exact(size_t *, size_t *, size_t *);
double qas_exact(size_t *t);

bool categorical_init(struct categorical_supp *, size_t, size_t);
struct categorical_res categorical_impl(struct categorical_supp *, uint8_t *, size_t *, size_t, size_t, enum categorical_flags);
void categorical_close(struct categorical_supp *);

bool maver_adj_init(struct maver_adj_supp *, size_t, size_t, size_t);
struct maver_adj_res maver_adj_impl(struct maver_adj_supp *, uint8_t *, size_t *, size_t, size_t, size_t, size_t, size_t, gsl_rng *, enum categorical_flags);
void maver_adj_close(struct maver_adj_supp *);


text:
/* multifit_nlinear/lm.c
 *
 * Copyright (C) 2014, 2015, 2016 Patrick Alken
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include <config.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_multifit_nlinear.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_permutation.h>

/*
 * This module contains an implementation of the Levenberg-Marquardt
 * algorithm for nonlinear optimization problems. This implementation
 * closely follows the following works:
 *
 * [1] H. B. Nielsen, K. Madsen, Introduction to Optimization and
 *     Data Fitting, Informatics and Mathematical Modeling,
 *     Technical University of Denmark (DTU), 2010.
 *
 * [2] J. J. More, The Levenberg-Marquardt Algorithm: Implementation
 *     and Theory, Lecture Notes in Mathematics, v630, 1978.
 */

typedef struct
{
  size_t n;           /* number of observations */
  size_t p;           /* number of parameters */
  gsl_vector *fvv;    /* D_v^2 f(x), size n */
  gsl_vector *vel;    /* geodesic velocity (standard LM step), size p */
  gsl_vector *acc;    /* geodesic acceleration, size p */
  gsl_vector *workp;  /* workspace, length p */
  gsl_vector *workn;  /* workspace, length n */

  int accel;          /* use geodesic acceleration? */

  /* tunable parameters */
  gsl_multifit_nlinear_parameters params;
} lm_state_t;

#include "common.c"

static void *lm_alloc (const int accel, const void * params, const size_t n, const size_t p);
static void *lm_alloc_noaccel (const void * params, const size_t n, const size_t p);
static void *lm_alloc_accel (const void * params, const size_t n, const size_t p);
static void lm_free(void *vstate);
static int lm_init(const void *vtrust_state, void *vstate);
static int lm_preloop(const void * vtrust_state, void * vstate);
static int lm_step(const void * vtrust_state, const double delta,
                   gsl_vector * dx, void * vstate);
static int lm_preduction(const void * vtrust_state, const gsl_vector * dx,
                         double * pred, void * vstate);

static void *
lm_alloc (const int accel, const void * params, const size_t n, const size_t p)
{
  const gsl_multifit_nlinear_parameters *mparams = (const gsl_multifit_nlinear_parameters *) params;
  lm_state_t *state;

  state = calloc(1, sizeof(lm_state_t));
  if (state == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate lm state", GSL_ENOMEM);
    }

  state->workp = gsl_vector_alloc(p);
  if (state->workp == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for workp", GSL_ENOMEM);
    }

  state->workn = gsl_vector_alloc(n);
  if (state->workn == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for workn", GSL_ENOMEM);
    }

  state->fvv = gsl_vector_alloc(n);
  if (state->fvv == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for fvv", GSL_ENOMEM);
    }

  state->vel = gsl_vector_alloc(p);
  if (state->vel == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for vel", GSL_ENOMEM);
    }

  state->acc = gsl_vector_alloc(p);
  if (state->acc == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for acc", GSL_ENOMEM);
    }

  state->n = n;
  state->p = p;
  state->params = *mparams;
  state->accel = accel;

  return state;
}

static void *
lm_alloc_noaccel (const void * params, const size_t n, const size_t p)
{
  return lm_alloc(0, params, n, p);
}

static void *
lm_alloc_accel (const void * params, const size_t n, const size_t p)
{
  return lm_alloc(1, params, n, p);
}

static void
lm_free(void *vstate)
{
  lm_state_t *state = (lm_state_t *) vstate;

  if (state->workp)
    gsl_vector_free(state->workp);

  if (state->workn)
    gsl_vector_free(state->workn);

  if (state->fvv)
    gsl_vector_free(state->fvv);

  if (state->vel)
    gsl_vector_free(state->vel);

  if (state->acc)
    gsl_vector_free(state->acc);

  free(state);
}

/*
lm_init()
  Initialize LM solver

Inputs: vtrust_state - trust state
        vstate       - workspace

Return: success/error
*/

static int
lm_init(const void *vtrust_state, void *vstate)
{
  const gsl_multifit_nlinear_trust_state *trust_state =
    (const gsl_multifit_nlinear_trust_state *) vtrust_state;
  lm_state_t *state = (lm_state_t *) vstate;

  gsl_vector_set_zero(state->vel);
  gsl_vector_set_zero(state->acc);

  *(trust_state->avratio) = 0.0;

  return GSL_SUCCESS;
}

/*
lm_preloop()
  Initialize LM method for new Jacobian matrix
*/

static int
lm_preloop(const void * vtrust_state, void * vstate)
{
  int status;
  const gsl_multifit_nlinear_trust_state *trust_state =
    (const gsl_multifit_nlinear_trust_state *) vtrust_state;
  const gsl_multifit_nlinear_parameters *params = trust_state->params;

  (void)vstate;

  /* initialize linear least squares solver */
  status = (params->solver->init)(trust_state, trust_state->solver_state);
  if (status)
    return status;

  return GSL_SUCCESS;
}

/*
lm_step()
  Calculate a new step vector by solving the linear
least squares system:

[      J     ] v = - [ f ]
[ sqrt(mu) D ]       [ 0 ]
*/

static int
lm_step(const void * vtrust_state, const double delta,
        gsl_vector * dx, void * vstate)
{
  int status;
  const gsl_multifit_nlinear_trust_state *trust_state =
    (const gsl_multifit_nlinear_trust_state *) vtrust_state;
  lm_state_t *state = (lm_state_t *) vstate;
  const gsl_multifit_nlinear_parameters *params = trust_state->params;
  const double mu = *(trust_state->mu);

  (void)delta;

  /* prepare the linear solver with current LM parameter mu */
  status = (params->solver->presolve)(mu, trust_state, trust_state->solver_state);
  if (status)
    return status;

  /*
   * solve: [     J      ] v = - [ f ]
   *        [ sqrt(mu)*D ]       [ 0 ]
   */
  status = (params->solver->solve)(trust_state->f,
                                   state->vel,
                                   trust_state,
                                   trust_state->solver_state);
  if (status)
    return status;

  if (state->accel)
    {
      double anorm, vnorm;

      /* compute geodesic acceleration */
      status = gsl_multifit_nlinear_eval_fvv(params->h_fvv,
                                             trust_state->x,
                                             state->vel,
                                             trust_state->f,
                                             trust_state->J,
                                             trust_state->sqrt_wts,
                                             trust_state->fdf,
                                             state->fvv,
                                             state->workp);
      if (status)
        return status;

      /*
       * solve: [     J      ] a = - [ fvv ]
       *        [ sqrt(mu)*D ]       [  0  ]
       */
      status = (params->solver->solve)(state->fvv,
                                       state->acc,
                                       trust_state,
                                       trust_state->solver_state);
      if (status)
        return status;

      anorm = gsl_blas_dnrm2(state->acc);
      vnorm = gsl_blas_dnrm2(state->vel);

      /* store |a| / |v| */
      *(trust_state->avratio) = anorm / vnorm;
    }

  /* compute step dx = v + 1/2 a */
  scaled_addition(1.0, state->vel, 0.5, state->acc, dx);

  return GSL_SUCCESS;
}

/*
lm_preduction()
  Compute predicted reduction using Eq 4.4 of More 1978
*/

static int
lm_preduction(const void * vtrust_state, const gsl_vector * dx,
              double * pred, void * vstate)
{
  const gsl_multifit_nlinear_trust_state *trust_state =
    (const gsl_multifit_nlinear_trust_state *) vtrust_state;
  lm_state_t *state = (lm_state_t *) vstate;
  const gsl_vector *diag = trust_state->diag;
  const gsl_vector *p = state->vel;
  const double norm_Dp = scaled_enorm(diag, p);
  const double normf = gsl_blas_dnrm2(trust_state->f);
  const double mu = *(trust_state->mu);
  double norm_Jp;
  double u, v;

  (void)dx;

  /* compute work = J*p */
  gsl_blas_dgemv(CblasNoTrans, 1.0, trust_state->J, p, 0.0, state->workn);

  /* compute ||J*p|| */
  norm_Jp = gsl_blas_dnrm2(state->workn);

  u = norm_Jp / normf;
  v = norm_Dp / normf;

  *pred = u * u + 2.0 * mu * v * v;

  return GSL_SUCCESS;
}

static const gsl_multifit_nlinear_trs lm_type =
{
  "levenberg-marquardt",
  lm_alloc_noaccel,
  lm_init,
  lm_preloop,
  lm_step,
  lm_preduction,
  lm_free
};

const gsl_multifit_nlinear_trs *gsl_multifit_nlinear_trs_lm = &lm_type;

static const gsl_multifit_nlinear_trs lmaccel_type =
{
  "levenberg-marquardt+accel",
  lm_alloc_accel,
  lm_init,
  lm_preloop,
  lm_step,
  lm_preduction,
  lm_free
};

const gsl_multifit_nlinear_trs *gsl_multifit_nlinear_trs_lmaccel = &lmaccel_type;

max_stars_repo_name: ielomariala/Hex-Game

max_issues_repo_name: Brian-ning/HMNE

max_forks_repo_name: Brian-ning/HMNE