Add cadnano vhelix and z_idx information via callbacks in new list-based reader

mrdna/readers/segmentmodel_from_cadnano.py

@@ -268,9 +268,9 @@ def mrdna_model_from_cadnano(json_file,seq=None,**model_parameters):
     stack=np.array((list(nt_prop["stack"])))
     orientation=np.array(list(nt_prop["orientation"]))
     rank = -np.array(list(nt_prop["z"])) # arrange segments in usual way
-    rank = rank + (np.max(rank)+1)*np.array(list(nt_prop["vh"]))
-    devlogger.info(f'Building model from lists')
+    _vh = np.array(list(nt_prop["vh"]))
 
+    devlogger.info(f'Building model from lists')
     model = model_from_basepair_stack_3prime( r, bp, stack, three_prime, seq, orientation,
                                               rank = rank, explicitly_unstacked = np.argwhere(stack == -1),
                                               **model_parameters )
@@ -282,14 +282,49 @@ def mrdna_model_from_cadnano(json_file,seq=None,**model_parameters):
     _starts = dict()
     for seg in model.segments:
         hid = int(seg._helix_idx)
-        sl = model._reader_list_hmap == hid
-        start_idx, end_idx = [np.argwhere(fwd & (hrank == i) & (hmap == hid)).flatten() for i in (0,seg.num_nt-1)]
+        _tmp = fwd & (hmap == hid) # select forward nts in helix
+        _list_ids = np.argwhere(_tmp).flatten()
+        _list_ids = _list_ids[np.argsort( nt_prop['zid'][_list_ids] )]
+
+        start_idx, end_idx = (_list_ids[0],_list_ids[-1])
         for _a in (start_idx, end_idx): assert( _a.size == 1 )
         start = max(map(float, [nt_prop['zid'][i] for i in (start_idx,end_idx)] ))
+
         cadnano_helix = int(nt_prop['vh'][start_idx])
         if cadnano_helix not in _starts: _starts[cadnano_helix] = []
         _starts[cadnano_helix].append( (seg,start) )
 
+        ## Assign beta and occupancy
+        seg._cadnano_bp_to_zidx = np.array( np.floor(nt_prop['zid'][_list_ids].astype(float)) )
+
+        seg.occupancy = cadnano_helix
+        def callback(segment):
+            for b in segment.beads:
+                bp = int(round(b.get_nt_position(segment)))
+                if bp < 0: bp = 0
+                if bp >= segment.num_nt: bp = segment.num_nt-1
+                try:
+                    b.beta = segment._cadnano_bp_to_zidx[bp]
+                    if 'orientation_bead' in b.__dict__:
+                        b.orientation_bead.beta = segment._cadnano_bp_to_zidx[bp]
+                except:
+                    pass
+        seg._generate_bead_callbacks.append(callback)
+
+        def atomic_callback(nucleotide, bp_to_zidx=seg._cadnano_bp_to_zidx):
+            nt = nucleotide
+            segment = nucleotide.parent.segment
+            bp = int(round(segment.contour_to_nt_pos( nt.contour_position )))
+            if bp < 0: bp = 0
+            if bp >= segment.num_nt: bp = segment.num_nt-1
+            try:
+                nt.beta = segment.bp_to_zidx[bp]
+                nt.parent.occupancy = segment.occupancy
+            except:
+                pass
+        seg._generate_nucleotide_callbacks.append(atomic_callback)
+
+
     for vh,ss in _starts.items():
         devlogger.debug(f'Renaming segs in {vh}')
         for i,(seg,s) in enumerate(sorted(ss, key=lambda x: x[-1])):