503 def Exec(self, opt) :
504 self.kalman_tracks.Clear('C')
505
506
507 if self.scale>1 and self.standalone:
508 if ROOT.gRandom.Rndm() > 1.0/self.scale: return
509
510 self.events_run += 1
511 hit_collection = {"pos" : [[], [], []],
512 "d" : [[], [], []],
513 "vert" : [],
514 "index" : [],
515 "system" : [],
516 "detectorID" : [],
517 "B" : [[], [], []],
518 "time": [],
519 "mask": []}
520
521 if ("us" in self.hits_to_fit) or ("ds" in self.hits_to_fit) or ("ve" in self.hits_to_fit) :
522
523 for i_hit, muFilterHit in enumerate(self.MuFilterHits) :
524
525 if muFilterHit.GetSystem() == 1 :
526 if "ve" not in self.hits_to_fit :
527 continue
528 elif muFilterHit.GetSystem() == 2 :
529 if "us" not in self.hits_to_fit :
530 continue
531 elif muFilterHit.GetSystem() == 3 :
532 if "ds" not in self.hits_to_fit :
533 continue
534 else :
535 if self.logger.IsLogNeeded(ROOT.fair.Severity.warn):
536 print("WARNING! Unknown MuFilter system!!")
537
538 self.mufiDet.GetPosition(muFilterHit.GetDetectorID(), self.a, self.b)
539
540 hit_collection["pos"][0].append(self.a.X())
541 hit_collection["pos"][1].append(self.a.Y())
542 hit_collection["pos"][2].append(self.a.Z())
543
544 hit_collection["B"][0].append(self.b.X())
545 hit_collection["B"][1].append(self.b.Y())
546 hit_collection["B"][2].append(self.b.Z())
547
548 hit_collection["vert"].append(muFilterHit.isVertical())
549 hit_collection["system"].append(muFilterHit.GetSystem())
550
551 hit_collection["d"][0].append(self.MuFilter_ds_dx)
552 hit_collection["d"][2].append(self.MuFilter_ds_dz)
553
554 hit_collection["index"].append(i_hit)
555
556 hit_collection["detectorID"].append(muFilterHit.GetDetectorID())
557 hit_collection["mask"].append(False)
558
559 times = []
560
561 if muFilterHit.GetSystem() == 3 :
562 hit_collection["d"][1].append(self.MuFilter_ds_dx)
563 for ch in range(self.MuFilter_ds_nSiPMs_hor):
564 if muFilterHit.isVertical() and ch==self.MuFilter_ds_nSiPMs_vert: break
565 if self.isMC:
566 times.append(muFilterHit.GetAllTimes()[ch])
567 else: times.append(muFilterHit.GetAllTimes()[ch]*6.25)
568
569 else :
570 hit_collection["d"][1].append(self.MuFilter_us_dy)
571 for ch in range(self.MuFilter_us_nSiPMs):
572 if self.isMC:
573 times.append(muFilterHit.GetAllTimes()[ch])
574 else: times.append(muFilterHit.GetAllTimes()[ch]*6.25)
575 hit_collection["time"].append(times)
576
577 if "sf" in self.hits_to_fit :
578 if self.Scifi_meas:
579
580 self.clusScifi.Clear()
581 self.scifiCluster()
582
583
584 for i_clust, scifiCl in enumerate(self.clusScifi) :
585 scifiCl.GetPosition(self.a,self.b)
586
587 hit_collection["pos"][0].append(self.a.X())
588 hit_collection["pos"][1].append(self.a.Y())
589 hit_collection["pos"][2].append(self.a.Z())
590
591 hit_collection["B"][0].append(self.b.X())
592 hit_collection["B"][1].append(self.b.Y())
593 hit_collection["B"][2].append(self.b.Z())
594
595
596 hit_collection["d"][0].append(scifiCl.GetN()*self.Scifi_dx)
597 hit_collection["d"][1].append(scifiCl.GetN()*self.Scifi_dy)
598 hit_collection["d"][2].append(self.Scifi_dz)
599
600 if int(scifiCl.GetFirst()/100000)%10==1:
601 hit_collection["vert"].append(True)
602 else: hit_collection["vert"].append(False)
603 hit_collection["index"].append(i_clust)
604
605 hit_collection["system"].append(0)
606 hit_collection["detectorID"].append(scifiCl.GetFirst())
607 hit_collection["mask"].append(False)
608
609 times = []
610 if self.isMC : times.append(scifiCl.GetTime()/6.25)
611 else: times.append(scifiCl.GetTime())
612 hit_collection["time"].append(times)
613
614 else:
615 if self.hits_for_triplet == 'sf' and self.hits_to_fit == 'sf':
616
617 N_plane_ZY = {1:0, 2:0, 3:0, 4:0, 5:0}
618 N_plane_ZX = {1:0, 2:0, 3:0, 4:0, 5:0}
619 for scifiHit in self.ScifiHits:
620 if not scifiHit.isValid(): continue
621 if scifiHit.isVertical():
622 N_plane_ZX[scifiHit.GetStation()] += 1
623 else:
624 N_plane_ZY[scifiHit.GetStation()] += 1
625 if self.mask_plane:
626 mask_plane_ZY = []
627 mask_plane_ZX = []
628
629 N_plane_ZY = dict(sorted(N_plane_ZY.items(), key=lambda item: item[1], reverse = True))
630 N_plane_ZX = dict(sorted(N_plane_ZX.items(), key=lambda item: item[1], reverse = True))
631
632 n_zx = self.Scifi_nPlanes - list(N_plane_ZX.values()).count(0)
633 n_zy = self.Scifi_nPlanes - list(N_plane_ZY.values()).count(0)
634
635 if n_zx < self.min_planes_hit or n_zy < self.min_planes_hit: return
636
637 for ii in range(n_zx-self.min_planes_hit):
638 if list(N_plane_ZX.values())[ii] > self.Nhits_per_plane:
639 mask_plane_ZX.append(list(N_plane_ZX.keys())[ii])
640 for ii in range(n_zy-self.min_planes_hit):
641 if list(N_plane_ZY.values())[ii] > self.Nhits_per_plane:
642 mask_plane_ZY.append(list(N_plane_ZY.keys())[ii])
643
644
645 for i_hit, scifiHit in enumerate(self.ScifiHits) :
646 if not scifiHit.isValid(): continue
647 self.scifiDet.GetSiPMPosition(scifiHit.GetDetectorID(), self.a, self.b)
648 hit_collection["pos"][0].append(self.a.X())
649 hit_collection["pos"][1].append(self.a.Y())
650 hit_collection["pos"][2].append(self.a.Z())
651
652 hit_collection["B"][0].append(self.b.X())
653 hit_collection["B"][1].append(self.b.Y())
654 hit_collection["B"][2].append(self.b.Z())
655
656 hit_collection["d"][0].append(self.Scifi_dx)
657 hit_collection["d"][1].append(self.Scifi_dy)
658 hit_collection["d"][2].append(self.Scifi_dz)
659
660 hit_collection["vert"].append(scifiHit.isVertical())
661 hit_collection["index"].append(i_hit)
662
663 hit_collection["system"].append(0)
664
665 hit_collection["detectorID"].append(scifiHit.GetDetectorID())
666
667 if self.hits_for_triplet == 'sf' and self.hits_to_fit == 'sf' and self.mask_plane:
668 if (scifiHit.isVertical()==0 and scifiHit.GetStation() in mask_plane_ZY) or (scifiHit.isVertical() and scifiHit.GetStation() in mask_plane_ZX):
669 hit_collection["mask"].append(True)
670 else: hit_collection["mask"].append(False)
671 else:
672 hit_collection["mask"].append(False)
673
674 times = []
675
676 if self.isMC : times.append(scifiHit.GetTime())
677 else: times.append(scifiHit.GetTime()*6.25)
678 hit_collection["time"].append(times)
679
680
681 if len(hit_collection['pos'][0])==0: return
682
683
684 for key, item in hit_collection.items() :
685 if key == 'vert' :
686 this_dtype = np.bool_
687 elif key == 'mask' :
688 this_dtype = np.bool_
689 elif key == "index" or key == "system" or key == "detectorID" :
690 this_dtype = np.int32
691 elif key != 'time' :
692 this_dtype = np.float32
693 if key== 'time':
694 length = max(map(len, item))
695 hit_collection[key] = np.stack(np.array([xi+[None]*(length-len(xi)) for xi in item]), axis = 1)
696 else: hit_collection[key] = np.array(item, dtype = this_dtype)
697
698
699 triplet_condition_system = []
700 if "sf" in self.hits_for_triplet :
701 triplet_condition_system.append(0)
702 if "ve" in self.hits_for_triplet :
703 triplet_condition_system.append(1)
704 if "us" in self.hits_for_triplet :
705 triplet_condition_system.append(2)
706 if "ds" in self.hits_for_triplet :
707 triplet_condition_system.append(3)
708
709
710 for i_muon in range(self.max_reco_muons) :
711
712 triplet_hits_horizontal = np.logical_and( ~hit_collection["vert"],
713 np.isin(hit_collection["system"], triplet_condition_system) )
714 triplet_hits_vertical = np.logical_and( hit_collection["vert"],
715 np.isin(hit_collection["system"], triplet_condition_system) )
716
717 n_planes_ZY = numPlanesHit(hit_collection["system"][triplet_hits_horizontal],
718 hit_collection["detectorID"][triplet_hits_horizontal])
719 if n_planes_ZY < self.min_planes_hit :
720 break
721
722 n_planes_ZX = numPlanesHit(hit_collection["system"][triplet_hits_vertical],
723 hit_collection["detectorID"][triplet_hits_vertical])
724 if n_planes_ZX < self.min_planes_hit :
725 break
726
727
728 muon_hits_horizontal = np.logical_and( np.logical_and( ~hit_collection["vert"], ~hit_collection["mask"]),
729 np.isin(hit_collection["system"], [1, 2, 3]))
730 muon_hits_vertical = np.logical_and( np.logical_and( hit_collection["vert"], ~hit_collection["mask"]),
731 np.isin(hit_collection["system"], [1, 2, 3]))
732 scifi_hits_horizontal = np.logical_and( np.logical_and( ~hit_collection["vert"], ~hit_collection["mask"]),
733 np.isin(hit_collection["system"], [0]))
734 scifi_hits_vertical = np.logical_and( np.logical_and( hit_collection["vert"], ~hit_collection["mask"]),
735 np.isin(hit_collection["system"], [0]))
736
737
738 ZY = np.dstack([np.concatenate([np.tile(hit_collection["pos"][2][muon_hits_horizontal], self.muon_weight),
739 hit_collection["pos"][2][scifi_hits_horizontal]]),
740 np.concatenate([np.tile(hit_collection["pos"][1][muon_hits_horizontal], self.muon_weight),
741 hit_collection["pos"][1][scifi_hits_horizontal]])])[0]
742
743 d_ZY = np.dstack([np.concatenate([np.tile(hit_collection["d"][2][muon_hits_horizontal], self.muon_weight),
744 hit_collection["d"][2][scifi_hits_horizontal]]),
745 np.concatenate([np.tile(hit_collection["d"][1][muon_hits_horizontal], self.muon_weight),
746 hit_collection["d"][1][scifi_hits_horizontal]])])[0]
747
748 ZX = np.dstack([np.concatenate([np.tile(hit_collection["pos"][2][muon_hits_vertical], self.muon_weight),
749 hit_collection["pos"][2][scifi_hits_vertical]]),
750 np.concatenate([np.tile(hit_collection["pos"][0][muon_hits_vertical], self.muon_weight),
751 hit_collection["pos"][0][scifi_hits_vertical]])])[0]
752
753 d_ZX = np.dstack([np.concatenate([np.tile(hit_collection["d"][2][muon_hits_vertical], self.muon_weight),
754 hit_collection["d"][2][scifi_hits_vertical]]),
755 np.concatenate([np.tile(hit_collection["d"][0][muon_hits_vertical], self.muon_weight),
756 hit_collection["d"][0][scifi_hits_vertical]])])[0]
757
758 is_scaled = False
759 ZY_hough = self.h_ZY.fit_randomize(ZY, d_ZY, self.n_random, is_scaled, self.draw)
760 ZX_hough = self.h_ZX.fit_randomize(ZX, d_ZX, self.n_random, is_scaled, self.draw)
761
762 tol = self.tolerance
763
764
765 if len(hit_collection["detectorID"]) <= self.max_n_Scifi_hits and self.hits_for_triplet == 'sf' and self.hits_to_fit == 'sf' :
766
767 if max(N_plane_ZX.values()) <= self.max_n_hits_plane and max(N_plane_ZY.values()) <= self.max_n_hits_plane :
768 tol = 5*self.tolerance
769
770 if self.hits_for_triplet == 'ds' and self.hits_to_fit == 'ds' :
771
772 N_plane_ZY = {0:0, 1:0, 2:0, 3:0}
773 N_plane_ZX = {0:0, 1:0, 2:0, 3:0}
774 for item in range(len(hit_collection["detectorID"])):
775 if hit_collection["vert"][item]: N_plane_ZX[(hit_collection["detectorID"][item]%10000)//1000] += 1
776 else: N_plane_ZY[(hit_collection["detectorID"][item]%10000)//1000] += 1
777 if max(N_plane_ZX.values()) <= self.max_n_hits_plane and max(N_plane_ZY.values()) <= self.max_n_hits_plane :
778 tol = 3*self.tolerance
779
780
781 track_hits_for_triplet_ZY = hit_finder(ZY_hough[0], ZY_hough[1],
782 np.dstack([hit_collection["pos"][2][triplet_hits_horizontal],
783 hit_collection["pos"][1][triplet_hits_horizontal]]),
784 np.dstack([hit_collection["d"][2][triplet_hits_horizontal],
785 hit_collection["d"][1][triplet_hits_horizontal]]), tol)
786
787 track_hits_for_triplet_ZX = hit_finder(ZX_hough[0], ZX_hough[1],
788 np.dstack([hit_collection["pos"][2][triplet_hits_vertical],
789 hit_collection["pos"][0][triplet_hits_vertical]]),
790 np.dstack([hit_collection["d"][2][triplet_hits_vertical],
791 hit_collection["d"][0][triplet_hits_vertical]]), tol)
792
793 n_planes_hit_ZY = numPlanesHit(hit_collection["system"][triplet_hits_horizontal][track_hits_for_triplet_ZY],
794 hit_collection["detectorID"][triplet_hits_horizontal][track_hits_for_triplet_ZY])
795
796 n_planes_hit_ZX = numPlanesHit(hit_collection["system"][triplet_hits_vertical][track_hits_for_triplet_ZX],
797 hit_collection["detectorID"][triplet_hits_vertical][track_hits_for_triplet_ZX])
798
799
800 if (self.hits_to_fit == 'sf' and len(hit_collection["detectorID"]) <= self.max_n_Scifi_hits and \
801 (n_planes_hit_ZY < self.min_planes_hit or n_planes_hit_ZX < self.min_planes_hit)):
802
803 is_scaled = True
804 ZY_hough = self.h_ZY.fit_randomize(ZY, d_ZY, self.n_random, is_scaled, self.draw)
805 ZX_hough = self.h_ZX.fit_randomize(ZX, d_ZX, self.n_random, is_scaled, self.draw)
806
807
808 track_hits_for_triplet_ZY = hit_finder(ZY_hough[0], ZY_hough[1],
809 np.dstack([hit_collection["pos"][2][triplet_hits_horizontal],
810 hit_collection["pos"][1][triplet_hits_horizontal]]),
811 np.dstack([hit_collection["d"][2][triplet_hits_horizontal],
812 hit_collection["d"][1][triplet_hits_horizontal]]), tol)
813
814 n_planes_hit_ZY = numPlanesHit(hit_collection["system"][triplet_hits_horizontal][track_hits_for_triplet_ZY],
815 hit_collection["detectorID"][triplet_hits_horizontal][track_hits_for_triplet_ZY])
816 if n_planes_hit_ZY < self.min_planes_hit:
817 break
818
819 track_hits_for_triplet_ZX = hit_finder(ZX_hough[0], ZX_hough[1],
820 np.dstack([hit_collection["pos"][2][triplet_hits_vertical],
821 hit_collection["pos"][0][triplet_hits_vertical]]),
822 np.dstack([hit_collection["d"][2][triplet_hits_vertical],
823 hit_collection["d"][0][triplet_hits_vertical]]), tol)
824 n_planes_hit_ZX = numPlanesHit(hit_collection["system"][triplet_hits_vertical][track_hits_for_triplet_ZX],
825 hit_collection["detectorID"][triplet_hits_vertical][track_hits_for_triplet_ZX])
826
827 if n_planes_hit_ZY < self.min_planes_hit or n_planes_hit_ZX < self.min_planes_hit:
828 break
829
830
831
832
833
834 track_hits_ZY = hit_finder(ZY_hough[0], ZY_hough[1],
835 np.dstack([hit_collection["pos"][2][~hit_collection["vert"]],
836 hit_collection["pos"][1][~hit_collection["vert"]]]),
837 np.dstack([hit_collection["d"][2][~hit_collection["vert"]],
838 hit_collection["d"][1][~hit_collection["vert"]]]), tol)
839
840 track_hits_ZX = hit_finder(ZX_hough[0], ZX_hough[1],
841 np.dstack([hit_collection["pos"][2][hit_collection["vert"]],
842 hit_collection["pos"][0][hit_collection["vert"]]]),
843 np.dstack([hit_collection["d"][2][hit_collection["vert"]],
844 hit_collection["d"][0][hit_collection["vert"]]]), tol)
845
846 posM = ROOT.TVector3(0, 0, 0.)
847 momM = ROOT.TVector3(0,0,100.)
848
849
850 covM = ROOT.TMatrixDSym(6)
851 if self.hits_to_fit.find('sf') >= 0:
852 res = self.kalman_sigmaScifi_spatial
853 if self.hits_to_fit == 'ds':
854 res = self.kalman_sigmaMufiDS_spatial
855 for i in range(3): covM[i][i] = res*res
856 for i in range(3,6): covM[i][i] = ROOT.TMath.Power(res / (4.*2.) / ROOT.TMath.Sqrt(3), 2)
857 rep = ROOT.genfit.RKTrackRep(13)
858
859
860 state = ROOT.genfit.MeasuredStateOnPlane(rep)
861 rep.setPosMomCov(state, posM, momM, covM)
862
863
864 seedState = ROOT.TVectorD(6)
865 seedCov = ROOT.TMatrixDSym(6)
866 rep.get6DStateCov(state, seedState, seedCov)
867
868 theTrack = ROOT.genfit.Track(rep, seedState, seedCov)
869
870
871 hit_z = np.concatenate([hit_collection["pos"][2][hit_collection["vert"]][track_hits_ZX],
872 hit_collection["pos"][2][~hit_collection["vert"]][track_hits_ZY]])
873
874 hit_A0 = np.concatenate([hit_collection["pos"][0][hit_collection["vert"]][track_hits_ZX],
875 hit_collection["pos"][0][~hit_collection["vert"]][track_hits_ZY]])
876
877 hit_A1 = np.concatenate([hit_collection["pos"][1][hit_collection["vert"]][track_hits_ZX],
878 hit_collection["pos"][1][~hit_collection["vert"]][track_hits_ZY]])
879
880 hit_B0 = np.concatenate([hit_collection["B"][0][hit_collection["vert"]][track_hits_ZX],
881 hit_collection["B"][0][~hit_collection["vert"]][track_hits_ZY]])
882
883 hit_B1 = np.concatenate([hit_collection["B"][1][hit_collection["vert"]][track_hits_ZX],
884 hit_collection["B"][1][~hit_collection["vert"]][track_hits_ZY]])
885
886 hit_B2 = np.concatenate([hit_collection["B"][2][hit_collection["vert"]][track_hits_ZX],
887 hit_collection["B"][2][~hit_collection["vert"]][track_hits_ZY]])
888
889 hit_detid = np.concatenate([hit_collection["detectorID"][hit_collection["vert"]][track_hits_ZX],
890 hit_collection["detectorID"][~hit_collection["vert"]][track_hits_ZY]])
891
892 kalman_spatial_sigma = np.concatenate([hit_collection["d"][0][hit_collection["vert"]][track_hits_ZX] / 12**0.5,
893 hit_collection["d"][1][~hit_collection["vert"]][track_hits_ZY] / 12**0.5])
894
895
896 kalman_max_dis = np.concatenate([((hit_collection["d"][0][hit_collection["vert"]][track_hits_ZX]/2.)**2 +
897 (hit_collection["d"][2][hit_collection["vert"]][track_hits_ZX]/2.)**2)**0.5,
898 ((hit_collection["d"][1][~hit_collection["vert"]][track_hits_ZY]/2.)**2 +
899 (hit_collection["d"][2][~hit_collection["vert"]][track_hits_ZY]/2.)**2)**0.5])
900
901 hitID = 0
902
903 hit_time = {}
904 for ch in range(hit_collection["time"].shape[0]):
905 hit_time[ch] = np.concatenate([hit_collection["time"][ch][hit_collection["vert"]][track_hits_ZX],
906 hit_collection["time"][ch][~hit_collection["vert"]][track_hits_ZY]])
907
908 for i_z_sorted in hit_z.argsort() :
909 tp = ROOT.genfit.TrackPoint()
910 hitCov = ROOT.TMatrixDSym(7)
911 hitCov[6][6] = kalman_spatial_sigma[i_z_sorted]**2
912
913 measurement = ROOT.genfit.WireMeasurement(ROOT.TVectorD(7, array('d', [hit_A0[i_z_sorted],
914 hit_A1[i_z_sorted],
915 hit_z[i_z_sorted],
916 hit_B0[i_z_sorted],
917 hit_B1[i_z_sorted],
918 hit_B2[i_z_sorted],
919 0.])),
920 hitCov,
921 1,
922 6,
923 tp)
924
925 measurement.setMaxDistance(kalman_max_dis[i_z_sorted])
926 measurement.setDetId(int(hit_detid[i_z_sorted]))
927 measurement.setHitId(int(hitID))
928 hitID += 1
929 tp.addRawMeasurement(measurement)
930 theTrack.insertPoint(tp)
931
932 if not theTrack.checkConsistency():
933 theTrack.Delete()
934 raise RuntimeError("Kalman fitter track consistency check failed.")
935
936
937 self.kalman_fitter.processTrack(theTrack)
938
939 fitStatus = theTrack.getFitStatus()
940 if not fitStatus.isFitConverged() and 0>1:
941 theTrack.Delete()
942 raise RuntimeError("Kalman fit did not converge.")
943
944
945 theTrack.SetUniqueID(self.track_type)
946 if fitStatus.isFitConverged():
947 if self.genfitTrack: self.kalman_tracks.Add(theTrack)
948 else :
949
950 this_track = ROOT.sndRecoTrack(theTrack)
951 pointTimes = ROOT.std.vector(ROOT.std.vector('float'))()
952 for n, i_z_sorted in enumerate(hit_z.argsort()):
953 t_per_hit = []
954 for ch in range(len(hit_time)):
955 if hit_time[ch][i_z_sorted] != None:
956 t_per_hit.append(hit_time[ch][i_z_sorted])
957 pointTimes.push_back(t_per_hit)
958 this_track.setRawMeasTimes(pointTimes)
959 this_track.setTrackType(self.track_type)
960
961 self.kalman_tracks[i_muon] = this_track
962
963 theTrack.Delete()
964
965
966
967 index_to_remove_ZX = np.where(np.in1d(hit_collection["detectorID"], hit_collection["detectorID"][hit_collection["vert"]][track_hits_ZX]))[0]
968 index_to_remove_ZY = np.where(np.in1d(hit_collection["detectorID"], hit_collection["detectorID"][~hit_collection["vert"]][track_hits_ZY]))[0]
969
970 index_to_remove = np.concatenate([index_to_remove_ZX, index_to_remove_ZY])
971
972
973 for key in hit_collection.keys() :
974 if len(hit_collection[key].shape) == 1 :
975 hit_collection[key] = np.delete(hit_collection[key], index_to_remove)
976 elif len(hit_collection[key].shape) == 2 :
977 hit_collection[key] = np.delete(hit_collection[key], index_to_remove, axis = 1)
978 else :
979 raise Exception("Wrong number of dimensions found when deleting hits in iterative muon identification algorithm.")
980