Fixed bug in saving variables

TenzinCHW · May 12, 2024 · 2829c63 · 2829c63
1 parent 8e22544
commit 2829c63
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/scripts/blanche_expt.jl b/scripts/blanche_expt.jl
@@ -47,11 +47,11 @@ function trainondata(data, maxiter, winsz, batchsize, arraycast)
     input = ordered_window(data, winsz)
     # converge on data
     output = MEFK.convergedynamics(model, input |> arraycast) |> Array
-    # TODO instead of returning input and output, return index of traindata and uniqueout
-    in2inds, ininds = uniqueind(traindata)
+    # instead of returning input and output, return index of traindata and uniqueout
     uniqueout = MEFK.convergedynamics(model, traindata |> arraycast) |> Array
+    ininds = uniqueinds(traindata, input)
     combout, comboutcnt = combine_counts(uniqueout, counts)
-    out2inds, outinds = uniqueind(combout)
+    outinds = uniqueinds(combout, output)
     cpumodel = modeltocpu(model)
     inspkcnt = sum(traindata, dims=2)
     comboutspkcnt = sum(combout, dims=2)
@@ -88,9 +88,9 @@ function runexperiment(binsz, winszs, maxiter, path, batchsize, arraycast, param
                 println("processing $saveloc")
                 model, ininds, uniqueinput, inspkcnt, counts, outinds, combout, comboutspkcnt, comboutcnt, losses =
                     trainondata(data_split[i], maxiter, winsz, batchsize, arraycast)
-                savedata["$i"] = Dict("ininds"=>ininds, "outinds"=>outinds, "input"=>uniqueinput, "output"=>combout, "net"=>model, "inspikecnt"=>inspkcnt, "incnt"=>counts, "outspikecnt"=>comboutspkcnt, "outcnt"=>comboutcnt, "loss"=>losses)
+                savedata["$i"] = Dict("ininds"=>ininds, "outinds"=>outinds, "input"=>uniqueinput, "output"=>combout, "net"=>model, "inspikecnt"=>inspkcnt, "incount"=>counts, "outspikecnt"=>comboutspkcnt, "outcount"=>comboutcnt, "loss"=>losses)
                 MAT.matwrite(joinpath(cdmdir, "input", "$(DrWatson.savename(params))_$(i).mat"),
-                             Dict("inspikecnt"=>inspkcnt, "incnt"=>counts))
+                             Dict("inspikecnt"=>inspkcnt, "counts"=>counts))
                 MAT.matwrite(joinpath(cdmdir, "complete", "$(DrWatson.savename(params))_$(i).mat"),
                              Dict("cells"=>model.n, "spike_counts"=>comboutspkcnt, "counts"=>comboutcnt))
             end
@@ -100,7 +100,7 @@ function runexperiment(binsz, winszs, maxiter, path, batchsize, arraycast, param
                 nulldata = generatebernoulli(data_split[i])
                 model, ininds, uniqueinput, inspkcnt, counts, outinds, combout, comboutspkcnt, comboutcnt, losses =
                     trainondata(nulldata, maxiter, winsz, batchsize, arraycast)
-                savenull["$i"] = Dict("ininds"=>ininds, "outinds"=>outinds, "input"=>uniqueinput, "output"=>combout, "net"=>model, "inspikecnt"=>inspkcnt, "incnt"=>counts, "outspikecnt"=>comboutspkcnt, "outcnt"=>comboutcnt, "loss"=>losses)
+                savenull["$i"] = Dict("ininds"=>ininds, "outinds"=>outinds, "input"=>uniqueinput, "output"=>combout, "net"=>model, "inspikecnt"=>inspkcnt, "incount"=>counts, "outspikecnt"=>comboutspkcnt, "outcount"=>comboutcnt, "loss"=>losses)
                 MAT.matwrite(joinpath(cdmdir, "null", "$(DrWatson.savename(params))_$(i).mat"),
                              Dict("cells"=>model.n, "spike_counts"=>comboutspkcnt, "counts"=>comboutcnt))
             end

diff --git a/scripts/calculate_entropy.jl b/scripts/calculate_entropy.jl
@@ -41,6 +41,17 @@ function uniqueind(data)
 end
 
 
+"""Get ordered array of uniquedata to indices of uniquedata that make up origdata.
+uniquedata might not be in same order as unique(origdata, dims=1)."""
+function uniqueinds(uniquedata, origdata)
+    inddict = Dict()
+    for i in 1:size(uniquedata)[1]
+        inddict[uniquedata[i, :]] = i
+    end
+    [inddict[origdata[i, :]] for i in 1:size(origdata)[1]]
+end
+
+
 uniquecounts(arr) = [count(==(e), arr) for e in unique(arr)]
 # TODO figure out why this is so slow compared to the window function in extract_blanche.jl
 function uniquecounts(data, dims::Int)