From b34cae89729ebfc2cd2ce71e7717e8d27c949a31 Mon Sep 17 00:00:00 2001
From: naotoY <naotino@gmail.com>
Date: Mon, 6 Feb 2023 15:57:52 +0900
Subject: [PATCH] Add print_time and check_alive functions to util.
print_time : output the date and time with messages to strandard output.
check_alive : detect the charge loss after aperture function is called.
---
.gitignore | 1 +
mbtrack2/utilities/check_alive.py | 40 +++++++++++++++++++++++++++++++
mbtrack2/utilities/print_time.py | 20 ++++++++++++++++
3 files changed, 61 insertions(+)
create mode 100644 mbtrack2/utilities/check_alive.py
create mode 100644 mbtrack2/utilities/print_time.py
diff --git a/.gitignore b/.gitignore
index 2ec607f..065def9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,3 +2,4 @@
*.ipynb_checkpoints*
test_*.py
*.hdf5
+bk_*.py
\ No newline at end of file
diff --git a/mbtrack2/utilities/check_alive.py b/mbtrack2/utilities/check_alive.py
new file mode 100644
index 0000000..a35f0b9
--- /dev/null
+++ b/mbtrack2/utilities/check_alive.py
@@ -0,0 +1,40 @@
+###################################################################################
+### from nytool/utilities to mbtrack2/utilities
+### 2023/02/06 written by Naoto Yamamoto, naotoy@post.kek.jp
+###
+### These functions are made to keep the similar usage as the original mbtrack.
+###
+###################################################################################
+
+def check_alive(turn,beam,filename=None):
+ """
+ turn : int
+ turn nuber
+ beam ; Beam or Bunch_object
+ beam or bunch object
+ I0 : float
+ total beam current
+ """
+ if check_alive.enable:
+ from mpi4py import MPI
+ comm = MPI.COMM_WORLD
+ rank = comm.Get_rank()
+ size = comm.Get_size()
+ if size == 1:
+ for idx in beam.bunch_index:
+ if False in beam[idx].alive:
+ print("!!!!! Beam loss at turn %d of Bunch %d !!!!!" % (turn,idx),flush=True)
+ if filename is not None:
+ with open("%s.loss.txt" %(filename),'a') as f:
+ print("!!!!! Beam loss at turn %d of Bunch %d !!!!!" % (turn,idx),file=f)
+ check_alive.enable = False
+ else: #MPI
+ #bunch = beam[beam.mpi.rank_to_bunch(size)]
+ bunch = beam[beam.mpi.bunch_num]
+ if False in bunch.alive:
+ print("!!!!! Beam loss at turn %d of Bunch %d !!!!!" % (turn,beam.mpi.bunch_num),flush=True)
+ if filename is not None:
+ with open("%s.loss.txt" %(filename),'a') as f:
+ print("!!!!! Beam loss at turn %d of Bunch %d !!!!!" % (turn,beam.mpi.bunch_num),file=f)
+ check_alive.enable = False
+check_alive.enable = True
diff --git a/mbtrack2/utilities/print_time.py b/mbtrack2/utilities/print_time.py
new file mode 100644
index 0000000..c8781f7
--- /dev/null
+++ b/mbtrack2/utilities/print_time.py
@@ -0,0 +1,20 @@
+###################################################################################
+### from nytool/utilities to mbtrack2/utilities
+### 2023/02/06 written by Naoto Yamamoto, naotoy@post.kek.jp
+###
+### These functions are made to keep the similar usage as the original mbtrack.
+###
+###################################################################################
+
+def print_time(message = "",flush=False):
+ """
+ Print a message only once even if mpi is used.
+ """
+
+ from time import strftime,localtime
+ from mpi4py import MPI
+ comm = MPI.COMM_WORLD
+ rank = comm.Get_rank()
+ if rank == 0:
+ mess = strftime(" [%d %b %Y %H:%M:%S] ", localtime()) + message
+ print(mess,flush=flush)
--
GitLab