gazzolalab / miv-simulator Goto Github PK

from miv_simulator.simulator import create_neural_h5

cell_distributions = {
  "STIM": {"SO": 0, "SP": 1000, "SR": 0, "SLM": 0},
  "PYR": {"SO": 0, "SP": 80000, "SR": 0, "SLM": 0},
  "PVBC": {"SO": 352, "SP": 1034, "SR": 88, "SLM": 0},
  "OLM": {"SO": 438, "SP": 0, "SR": 0, "SLM": 0},
}

create_neural_h5(
  output_filepath='./example.h5',
  cell_distributions=cell_distributions,
  synapses={
      'PYR': {'STIM': True, 'PYR': True, "PVBC": True, "OLM": True},
      'PVBC': {'PYR': True, 'STIM': True, 'PVBC': True, 'OLM': True},
      'OLM': {'PYR': True, 'PVBC': True}
  }
)

from miv_simulator.simulator import generate_network_architecture

layer_extents={
  "SO": [[0, 0, 0], [4000, 4000, 100]],
  "SP": [[0, 0, 100], [4000, 4000, 150]],
  "SR": [[0, 0, 150], [4000, 4000, 350]],
  "SLM": [[0, 0, 350], [4000, 4000, 450]],
}

rotation=[0, 0, 0]

generate_network_architecture(
  output_filepath="example.h5",
  cell_distributions=cell_distributions,
  layer_extents=layer_extents,
  rotation=rotation,
  cell_constraints={"PC": {"SP": [100, 120]}, "PVBC": {"SR": [150, 200]}},
  output_namespace= "Generated Coordinates",
  geometry_filepath=None,
  populations=(),
  resolution= (3, 3, 3),
  alpha_radius=2500.0,
  nodeiter=10,
  dispersion_delta=0.1,
  snap_delta=0.01,
  h5_types_filepath=None,
  io_size=-1,
  chunk_size=1000,
  value_chunk_size=1000,
)

import numpy as np
import matplotlib.cm as cm
import matplotlib.pyplot as plt
from miv_simulator.geometry.geometry import get_total_extents
from neuroh5.io import read_cell_attributes

# options
populations = ('STIM',)
mayavi = False # set to True for 3D
subvol = False
subpopulation = -1
scale = 25.0

print("Reading coordinates...")

(extent_u, extent_v, extent_l) = get_total_extents(layer_extents)

pop_min_extent = None
pop_max_extent = None

xcoords = []
ycoords = []
zcoords = []
cmap = cm.get_cmap("Dark2")
cmap_range = np.linspace(0, 1, num=len(populations))

colors = []
for pop_id, population in enumerate(populations):
  coords = read_cell_attributes(
      'example.h5', population, namespace="Generated Coordinates"
  )

  count = 0
  cxcoords = []
  cycoords = []
  czcoords = []
  for k, v in coords:
      count += 1
      cxcoords.append(v["X Coordinate"][0])
      cycoords.append(v["Y Coordinate"][0])
      czcoords.append(v["Z Coordinate"][0])
  if subpopulation > -1 and count > subpopulation:
      ridxs = np.random.choice(
          np.arange(len(cxcoords)), replace=False, size=subpopulation
      )
      cxcoords = list(np.asarray(cxcoords)[ridxs])
      cycoords = list(np.asarray(cycoords)[ridxs])
      czcoords = list(np.asarray(czcoords)[ridxs])

  colors += [cmap(cmap_range[pop_id]) for _ in range(len(cxcoords))]
  xcoords += cxcoords
  ycoords += cycoords
  zcoords += czcoords
  
  print(f"Read {count} coordinates...")

  pop_distribution = cell_distributions[population]
  pop_layers = []
  for layer in pop_distribution:
      num_layer = pop_distribution[layer]
      if num_layer > 0:
          pop_layers.append(layer)

          if pop_min_extent is None:
              pop_min_extent = np.asarray(layer_extents[layer][0])
          else:
              pop_min_extent = np.minimum(
                  pop_min_extent, np.asarray(layer_extents[layer][0])
              )

          if pop_max_extent is None:
              pop_max_extent = np.asarray(layer_extents[layer][1])
          else:
              pop_max_extent = np.maximum(
                  pop_min_extent, np.asarray(layer_extents[layer][1])
              )

pts = np.concatenate(
  (
      np.asarray(xcoords).reshape(-1, 1),
      np.asarray(ycoords).reshape(-1, 1),
      np.asarray(zcoords).reshape(-1, 1),
  ),
  axis=1,
)

if mayavi:
  from mayavi import mlab
else:
  import matplotlib.pyplot as plt

print("Plotting coordinates...")
if mayavi:
  mlab.points3d(*pts.T, color=(1, 1, 0), scale_factor=scale)
else:
  fig = plt.figure()
  ax = fig.add_subplot(111, projection="3d")
  ax.scatter(*pts.T, c=colors, s=int(scale))

print("Constructing volume...")
from miv_simulator.volume import make_network_volume

if subvol:
  subvol = make_network_volume(
      (pop_min_extent[0], pop_max_extent[0]),
      (pop_min_extent[1], pop_max_extent[1]),
      (pop_min_extent[2], pop_max_extent[2]),
      resolution=[3, 3, 3],
      rotate=rotation,
  )
else:
  vol = make_network_volume(
      (extent_u[0], extent_u[1]),
      (extent_v[0], extent_v[1]),
      (extent_l[0], extent_l[1]),
      resolution=[3, 3, 3],
      rotate=rotation,
  )

print("Plotting volume...")

if subvol:
  if mayavi:
      subvol.mplot_surface(color=(0, 0.4, 0), opacity=0.33)
  else:
      subvol.mplot_surface(color="k", alpha=0.33, figax=[fig, ax])
else:
  if mayavi:
      vol.mplot_surface(color=(0, 1, 0), opacity=0.33)
  else:
      vol.mplot_surface(color="k", alpha=0.33, figax=[fig, ax])
if mayavi:
  mlab.show()
else:
  ax.view_init(-90, 0)
  plt.show()