main.cpp

/*  main.cpp – minimal LIF simulator for the adult fly connectome
    Build :  make
    Run   :  ./main --csv connectome.csv --active act.txt --silent off.txt
             --t 1000
*/
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <fstream>
#include <iostream>
#include <numeric>
#include <random>
#include <string>
#include <tuple>
#include <vector>
#ifdef USE_ARROW
#include <arrow/api.h>
#include <arrow/io/api.h>
#include <parquet/arrow/reader.h>
#endif

using namespace std;

/* ---------- parameter block ---------- */
struct Params {
  float dt = 1.0f; // ms
  float v0 = -52.f, v_rst = -52.f, v_th = -45.f;
  float tau = 5.f;       // ms
  float t_mbr = 20.f;    // ms
  float t_rfc = 2.2f;    // ms
  float w_syn = .275f;   // mV
  float delay_ms = 1.0f; // axonal delay
  /* external drive */
  float poi_rate = 20.f;     // Hz
  float w_poi_scale = 250.f; // multiplied with w_syn
};
/* ---------- CSR graph ---------- */
struct CSR {
  vector<uint32_t> row; // size N+1
  vector<uint32_t> col;
  vector<float> w;
};
/* ---------- utility ---------- */
static vector<uint32_t> load_list(const string &path) {
  if (path.empty())
    return {};
  ifstream f(path);
  uint32_t id;
  vector<uint32_t> v;
  while (f >> id)
    v.push_back(id);
  return v;
}
static CSR load_csv_to_csr(const string &csv, size_t N,
                           const vector<uint8_t> &silent) {
  ifstream f(csv);
  string line;
  CSR g;
  vector<tuple<uint32_t, uint32_t, float>> edges;
  edges.reserve(54'000'000); // pre-reserve
  getline(f, line);          // skip header
  while (getline(f, line)) {
    uint32_t pre, post;
    float w;
    sscanf(line.c_str(), "%u,%u,%f", &pre, &post, &w);
    if (silent[pre] || silent[post])
      continue;
    edges.emplace_back(pre, post, w);
  }
  sort(edges.begin(), edges.end(),
       [](auto &a, auto &b) { return get<0>(a) < get<0>(b); });
  g.row.assign(N + 1, 0);
  g.col.reserve(edges.size());
  g.w.reserve(edges.size());
  for (auto [pre, post, w] : edges) {
    ++g.row[pre + 1];
    g.col.push_back(post);
    g.w.push_back(w);
  }
  partial_sum(g.row.begin(), g.row.end(), g.row.begin());
  return g;
}
#ifdef USE_ARROW
static size_t count_neurons_parquet(const string &pq) {
  std::shared_ptr<arrow::io::ReadableFile> infile;
  PARQUET_ASSIGN_OR_THROW(infile, arrow::io::ReadableFile::Open(pq));
  std::unique_ptr<parquet::arrow::FileReader> reader;
  PARQUET_THROW_NOT_OK(
      parquet::arrow::OpenFile(infile, arrow::default_memory_pool(), &reader));
  std::shared_ptr<arrow::Table> table;
  PARQUET_THROW_NOT_OK(
      reader->ReadTable({"Presynaptic_Index", "Postsynaptic_Index"}, &table));
  auto pre_chunks = table->GetColumnByName("Presynaptic_Index")->chunks();
  auto post_chunks = table->GetColumnByName("Postsynaptic_Index")->chunks();
  size_t N = 0;
  for (size_t c = 0; c < pre_chunks.size(); ++c) {
    auto pre = std::static_pointer_cast<arrow::UInt32Array>(pre_chunks[c]);
    auto post = std::static_pointer_cast<arrow::UInt32Array>(post_chunks[c]);
    for (int64_t i = 0; i < pre->length(); ++i) {
      uint32_t a = pre->Value(i);
      uint32_t b = post->Value(i);
      N = max(N, size_t(max(a, b) + 1));
    }
  }
  return N;
}
static CSR load_parquet_to_csr(const string &pq, size_t N,
                               const vector<uint8_t> &silent) {
  std::shared_ptr<arrow::io::ReadableFile> infile;
  PARQUET_ASSIGN_OR_THROW(infile, arrow::io::ReadableFile::Open(pq));
  std::unique_ptr<parquet::arrow::FileReader> reader;
  PARQUET_THROW_NOT_OK(
      parquet::arrow::OpenFile(infile, arrow::default_memory_pool(), &reader));
  std::shared_ptr<arrow::Table> table;
  PARQUET_THROW_NOT_OK(reader->ReadTable(&table));
  auto pre_chunks = table->GetColumnByName("Presynaptic_Index")->chunks();
  auto post_chunks = table->GetColumnByName("Postsynaptic_Index")->chunks();
  auto w_chunks = table->GetColumnByName("Connectivity")->chunks();
  vector<tuple<uint32_t, uint32_t, float>> edges;
  size_t total = table->num_rows();
  edges.reserve(total);
  for (size_t c = 0; c < pre_chunks.size(); ++c) {
    auto pre = std::static_pointer_cast<arrow::UInt32Array>(pre_chunks[c]);
    auto post = std::static_pointer_cast<arrow::UInt32Array>(post_chunks[c]);
    auto w = std::static_pointer_cast<arrow::FloatArray>(w_chunks[c]);
    for (int64_t i = 0; i < pre->length(); ++i) {
      uint32_t a = pre->Value(i);
      uint32_t b = post->Value(i);
      if (silent[a] || silent[b])
        continue;
      float wt = w->Value(i);
      edges.emplace_back(a, b, wt);
    }
  }
  sort(edges.begin(), edges.end(),
       [](auto &x, auto &y) { return get<0>(x) < get<0>(y); });
  CSR g;
  g.row.assign(N + 1, 0);
  g.col.reserve(edges.size());
  g.w.reserve(edges.size());
  for (auto [pre, post, w] : edges) {
    ++g.row[pre + 1];
    g.col.push_back(post);
    g.w.push_back(w);
  }
  partial_sum(g.row.begin(), g.row.end(), g.row.begin());
  return g;
}
#endif
/* ---------- simulation ---------- */
struct Simulator {
  size_t N;
  CSR G;
  Params P;
  vector<float> v, g; // per-neuron state
  vector<int> refr;   // refractory counter (steps)
  vector<uint8_t> silent;
  size_t delay_steps;
  struct Event {
    uint32_t id;
    float w;
  };
  vector<vector<Event>> ring;
  vector<uint64_t> spikes;

  Simulator(size_t n, const CSR &csr, const Params &p,
            vector<uint8_t> silent_neurons)
      : N(n), G(csr), P(p), v(n, p.v0), g(n, 0.f), refr(n, int(p.t_rfc / p.dt)),
        silent(std::move(silent_neurons)) {
    delay_steps = size_t(round(P.delay_ms / P.dt));
    ring.assign(delay_steps, {});
  }
  void drive_poisson(const vector<uint32_t> &active, size_t T) {
    std::mt19937_64 rng(42);
    std::exponential_distribution<float> exp(P.poi_rate / 1000.f); // per ms
    for (auto id : active) {
      if (silent[id])
        continue;
      float t = 0.f;
      while (t < T) {
        t += exp(rng);
        if (t < T) {
          size_t step = size_t(t);
          ring[step % delay_steps].push_back(
              Event{id, P.w_syn * P.w_poi_scale});
        }
      }
    }
  }
  void run(size_t T, const vector<uint32_t> &active) {
    drive_poisson(active, T);
    for (size_t step = 0; step < T; ++step) {
      /* 1. deliver queued events */
      for (const auto &ev : ring[step % delay_steps])
        g[ev.id] += ev.w;
      ring[step % delay_steps].clear();

      /* 2. membrane update (parallel) */
#pragma omp parallel for schedule(static)
      for (size_t i = 0; i < N; ++i) {
        if (silent[i])
          continue;
        if (refr[i] > 0) {
          --refr[i];
          continue;
        }
        float dv = (P.v0 - v[i] + g[i]) / P.t_mbr;
        v[i] += P.dt * dv;
        g[i] -= P.dt * g[i] / P.tau;
      }

      /* 3. spike detection (serial – rare) */
      for (size_t i = 0; i < N; ++i) {
        if (!silent[i] && refr[i] == 0 && v[i] > P.v_th) {
          // record
          spikes.push_back((uint64_t(i) << 32) | uint64_t(step));
          // schedule synaptic events
          for (uint32_t e = G.row[i]; e < G.row[i + 1]; ++e) {
            uint32_t j = G.col[e];
            float w = G.w[e];
            ring[(step + delay_steps) % delay_steps].push_back(Event{j, w});
          }
          // reset
          v[i] = P.v_rst;
          g[i] = 0.f;
          refr[i] = int(P.t_rfc / P.dt);
        }
      }
#ifdef VERIFY
      if (step % 100 == 0) {
        cerr << "step " << step << "\r";
      }
#endif
    }
  }
  void save_bin(const string &path) {
    ofstream o(path, ios::binary);
    o.write(reinterpret_cast<char *>(spikes.data()),
            spikes.size() * sizeof(uint64_t));
  }
};
/* ---------- main ---------- */
size_t run_simulation(const string &csv, const string &pq, const string &act,
                      const string &sil, size_t T) {
  size_t N = 0;
  CSR G;
#ifdef USE_ARROW
  if (!pq.empty()) {
    N = count_neurons_parquet(pq);
    vector<uint8_t> silent_vec_tmp(N, 0);
    for (auto id : load_list(sil))
      if (id < N)
        silent_vec_tmp[id] = 1;
    G = load_parquet_to_csr(pq, N, silent_vec_tmp);
    Params P;
    Simulator S(N, G, P, std::move(silent_vec_tmp));
    S.run(T, load_list(act));
    S.save_bin("spikes.bin");
    return S.spikes.size();
  }
#endif
  /* default CSV loader */
  {
    ifstream f(csv);
    string l;
    getline(f, l);
    while (getline(f, l)) {
      uint32_t pre, post;
      sscanf(l.c_str(), "%u,%u", &pre, &post);
      N = max(N, size_t(max(pre, post) + 1));
    }
  }
  vector<uint8_t> silent_vec(N, 0);
  for (auto id : load_list(sil))
    if (id < N)
      silent_vec[id] = 1;

  G = load_csv_to_csr(csv, N, silent_vec);
  Params P;
  Simulator S(N, G, P, silent_vec);
  S.run(T, load_list(act));
  S.save_bin("spikes.bin");
  return S.spikes.size();
}

extern "C" size_t run_simulation_c(const char *csv, const char *pq,
                                   const char *act, const char *sil, size_t T) {
  return run_simulation(csv ? string(csv) : "", pq ? string(pq) : "",
                        act ? string(act) : "", sil ? string(sil) : "", T);
}

#ifndef BENCH_LIB
int main(int argc, char **argv) {
  string csv, pq, act, sil;
  size_t T = 1000;
  for (int i = 1; i < argc; ++i) {
    string a = argv[i];
    if (a == "--csv")
      csv = argv[++i];
    else if (a == "--parquet")
      pq = argv[++i];
    else if (a == "--active")
      act = argv[++i];
    else if (a == "--silent")
      sil = argv[++i];
    else if (a == "--t")
      T = stoul(argv[++i]);
  }
  if (csv.empty() && pq.empty()) {
    cerr << "--csv or --parquet required\n";
    return 1;
  }
  size_t spikes = run_simulation(csv, pq, act, sil, T);
  cerr << "Spikes: " << spikes << "\n";
}
#endif