sieic_vis

#!/usr/bin/env python2

import argparse
from array import array
import math

def run(args):
    import ROOT, DDG4

    kernel = DDG4.Kernel()
    kernel.loadGeometry('file:' + args.compact_path)
    dd = kernel.detectorDescription()
    manager = dd.manager()
    if not args.nocheck:
        manager.CheckOverlaps(0.01)
        manager.PrintOverlaps()
    if args.input_file == None:
        manager.SetVisDensity()
    manager.SetVisLevel(2)
    top_vol = manager.GetTopVolume()
    top_vol.Draw('ogl')

    viewer = ROOT.gPad.GetViewer3D()
    viewer.SetCurrentCamera(ROOT.TGLViewer.kCameraOrthoZOY)
    viewer.ToggleOrthoRotate()
    if args.dark_theme:
        viewer.SwitchColorSet()

    if args.input_file == None:
        overlay = viewer.GetCameraOverlay()
        overlay.SetOrthographicMode(ROOT.TGLCameraOverlay.kAxis)
        overlay.SetShowOrthographic(True)

        clip_set = viewer.GetClipSet()
        clip_type = ROOT.TGLClip.kClipPlane
        clip_set.SetClipType(clip_type)
        clip_state = array('d')
        clip_state.extend([1,0,0,0,0,0])
        clip_set.SetClipState(clip_type, clip_state)
    else:
        markers, lines = draw_event(args.input_file, args.input_event)
        for marker in markers:
            marker.Draw('same')
        for line in lines:
            line.Draw('same')
        args.output_prefix += '_' + args.input_file + '_' + str(args.input_event)

    viewer.UpdateScene()

    if args.interactive:
        ROOT.gApplication.Run()
    else:
        save_picture(viewer, args.output_prefix + '_perp.png')
        camera = viewer.CurrentCamera()
        camera.RotateRad(0, ROOT.TMath.Pi() / 2.)
        save_picture(viewer, args.output_prefix + '_long.png')
        camera.RotateRad(0, -ROOT.TMath.Pi() / 2.)
        camera.RotateRad(-0.25, 0.4)
        save_picture(viewer, args.output_prefix + '_angled.png')
        if args.input_file == None:
            overlay.SetShowOrthographic(False)
            save_picture(viewer, args.output_prefix + '_angledNoOverlay.png')

def draw_event(input_file, input_event):
    import ROOT

    ROOT.gSystem.Load('liblcio')
    ROOT.gInterpreter.ProcessLine("#include <lcio.h>")
    ROOT.gInterpreter.ProcessLine("#include <EVENT/LCCollection.h>")
    ROOT.gInterpreter.ProcessLine("#include <IMPL/SimTrackerHitImpl.h>")
    ROOT.gInterpreter.ProcessLine("#include <IMPL/SimCalorimeterHitImpl.h>")
    ROOT.gInterpreter.ProcessLine("#include <IMPL/MCParticleImpl.h>")

    reader = ROOT.IOIMPL.LCFactory.getInstance().createLCReader()
    reader.open(input_file)
    reader.skipNEvents(input_event)
    event = reader.readNextEvent()

    markers = []
    lines = []

    coll_names = event.getCollectionNames()
    for coll_name in coll_names:
        coll = event.getCollection(coll_name)
        type_name = coll.getTypeName()
        if type_name == "SimTrackerHit":
            n_hits = coll.getNumberOfElements()
            tracker_hit_markers = ROOT.TPolyMarker3D()
            markers.append(tracker_hit_markers)
            tracker_hit_markers.SetMarkerSize(10)
            tracker_hit_markers.SetMarkerColor(3)
            tracker_hit_markers.SetMarkerStyle(2)
            for i in range(0, n_hits):
                hit = coll.getElementAt(i)
                hit.__class__ = ROOT.IMPL.SimTrackerHitImpl
                energy = hit.getEDep()
                if energy == 0:
                    continue
                pos = hit.getPosition()
                tracker_hit_markers.SetNextPoint(pos[0] / 10., pos[1] / 10., pos[2] / 10.)
        elif type_name == "SimCalorimeterHit":
            n_hits = coll.getNumberOfElements()
            for i in range(0, n_hits):
                hit = coll.getElementAt(i)
                hit.__class__ = ROOT.IMPL.SimCalorimeterHitImpl
                energy = hit.getEnergy()
                if energy == 0:
                    continue
                marker_size = math.log(energy * 1e12) * 0.25
                if(marker_size <= 0):
                    continue
                pos = hit.getPosition()
                calo_hit_markers = ROOT.TPolyMarker3D(1)
                markers.append(calo_hit_markers)
                calo_hit_markers.SetMarkerSize(marker_size)
                calo_hit_markers.SetMarkerColor(6)
                if coll_name.startswith("Hcal"):
                    calo_hit_markers.SetMarkerColor(7)
                calo_hit_markers.SetMarkerStyle(4)
                calo_hit_markers.SetPoint(0, pos[0] / 10., pos[1] / 10., pos[2] / 10.)
        elif type_name == "MCParticle":
            n_parts = coll.getNumberOfElements()
            for i in range(0, n_parts):
                part = coll.getElementAt(i)
                part.__class__ = ROOT.IMPL.MCParticleImpl
                if part.getGeneratorStatus() == 1:
                    vertex = part.getVertex()
                    vertex_marker = ROOT.TPolyMarker3D(1)
                    markers.append(vertex_marker)
                    vertex_marker.SetMarkerColor(3)
                    vertex_marker.SetMarkerStyle(7)
                    vertex_marker.SetPoint(0, vertex[0] / 10., vertex[1] / 10., vertex[2] / 10.)
                    vertex_marker.Draw('same')
                    momentum = part.getMomentum()
                    mag = math.sqrt(pow(momentum[0], 2) + pow(momentum[1], 2) + pow(momentum[2], 2))
                    line_len = 50
                    line_vec = array('d')
                    line_vec.extend([momentum[0] * line_len / mag, momentum[1] * line_len / mag, momentum[2] * line_len / mag])
                    part_line = ROOT.TPolyLine3D(2)
                    lines.append(part_line)
                    part_line.SetLineColor(1)
                    part_line.SetLineWidth(2)
                    part_line.SetPoint(0, vertex[0] / 10., vertex[1] / 10., vertex[2] / 10.)
                    part_line.SetPoint(1, line_vec[0], line_vec[1], line_vec[2])
                    part_line.Draw('same')

    reader.close()
    return markers, lines

def save_picture(viewer, filename):
    viewer.SavePictureUsingFBO(filename, 1360, 720)
    print("Saved file: " + filename)

def define_args(parser):
    parser.add_argument(
            'compact_path',
            metavar='CompactPath',
            help='file with compact XML detector description'
            )
    parser.add_argument(
            '-o', '--outputprefix',
            dest='output_prefix',
            default='sieic_vis',
            help='output file prefix'
            )
    parser.add_argument(
            '-i', '--interactive',
            action='store_true',
            dest='interactive',
            help='run an interactive session.  Must quit ROOT application from the menu.'
            )
    parser.add_argument(
            '-f', '--file',
            dest='input_file',
            default=None,
            help='lcio file path containing event to visualize'
            )
    parser.add_argument(
            '-e', '--event',
            type=int,
            dest='input_event',
            default=0,
            help='event to visualize'
            )
    parser.add_argument(
            '-n', '--nocheck',
            action='store_true',
            dest='nocheck',
            help='do not check geometry for overlaps'
            )
    parser.add_argument(
            '-d', '--dark',
            action='store_true',
            dest='dark_theme',
            help='use a black background'
            )

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='sieic_vis')
    define_args(parser)
    args = parser.parse_args()
    print(args)
    run(args)