def isJVMMemoryHigh():
url = "{0}/_nodes/stats?pretty=true".format(
random.choice(options.esservers))
r = requests.get(url)
logger.debug(r)
if r.status_code == 200:
nodestats = r.json()
for node in nodestats['nodes']:
loadaverage = nodestats['nodes'][node]['os']['cpu']['load_average']
cpuusage = nodestats['nodes'][node]['os']['cpu']['percent']
nodename = nodestats['nodes'][node]['name']
jvmused = nodestats['nodes'][node]['jvm']['mem'][
'heap_used_percent']
logger.debug('{0}: cpu {1}% jvm {2}% load average: {3}'.format(
nodename, cpuusage, jvmused, loadaverage))
if jvmused > options.jvmlimit:
logger.info(
'{0}: cpu {1}% jvm {2}% load average: {3} recommending cache clear'
.format(nodename, cpuusage, jvmused, loadaverage))
return True
return False
else:
logger.error(r)
return False
def apply_self_calibration(self, selfcal_config, calibration_mode,
output_ms_path, output_path, spw):
logger.info(
Color.HEADER +
"Applying self calibration for {0}".format(self._dataset_path) +
Color.ENDC)
cal_mode = selfcal_config['calmode']
mask_path = selfcal_config['masking']['mask_path'] if selfcal_config[
'masking']['mask_path'] else 'None'
script_path = 'casa_scripts/self_calibration.py'
script_parameters = "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10} {11} {12} {13} " \
"{14} {15}".format(config.CONFIG_PATH,
self._dataset_path,
output_path,
output_ms_path,
cal_mode[calibration_mode]['solint'],
config.GLOBAL_CONFIGS['refant'],
selfcal_config['minsnr'],
self._output_path,
cal_mode[calibration_mode]['applymode'],
selfcal_config['masking']['threshold'],
selfcal_config['masking']['bmask'],
mask_path,
cal_mode['ap']['loop_count'],
cal_mode['p']['loop_count'],
calibration_mode,
spw)
self._run(script_path, script_parameters)
def main():
if options.checkjvmmemory:
if isJVMMemoryHigh():
logger.info('initiating cache clearing')
clearESCache()
else:
clearESCache()
def extend_continuum_flags(self):
logger.info(Color.HEADER + "Extending continuum flags on line..." +
Color.ENDC)
flag_reasons = "{0},{1}".format(BAD_ANTENNA_TIME, BAD_BASELINE_TIME)
flag_file = "{0}/flags_continuum.txt".format(config.OUTPUT_PATH)
if os.path.exists(flag_file):
self.flagdata(flag_file, flag_reasons)
def analyse_antennas(self, spw_polarization_and_scan_product):
logger.info(Color.HEADER + "Started detailed flagging on all unflagged antennas" + Color.ENDC)
bad_window_present = False
for spw, polarization, scan_id in spw_polarization_and_scan_product:
scan_times = self.measurement_set.timesforscan(scan_id)
amp_matrix = AmplitudeMatrix(self.measurement_set, polarization, scan_id, spw, self._source_config)
global_median = amp_matrix.median()
global_sigma = amp_matrix.mad_sigma()
self._print_polarization_details(global_sigma, global_median, polarization, scan_id)
antennaids = self.measurement_set.antenna_ids(polarization, scan_id)
window_config = WindowConfig(*self._source_config['detail_flagging']['antenna_sliding_window'])
# Sliding Window for Bad Antennas
for antenna in antennaids:
filtered_matrix = amp_matrix.filter_by_antenna(antenna)
if filtered_matrix.has_sufficient_data(window_config) and filtered_matrix.is_bad(global_median, window_config.mad_scale_factor * global_sigma):
logger.info(
Color.FAIL + 'Antenna ' + str(
antenna) + ' is Bad running sliding Window on it' + Color.ENDC)
flagged_bad_window = self._flag_bad_time_window(BAD_ANTENNA_TIME, antenna,
filtered_matrix.amplitude_data_matrix,
global_sigma,
global_median, scan_times, polarization, scan_id,
window_config)
if flagged_bad_window: bad_window_present = True
return bad_window_present
def main():
# connect and declare the message queue/kombu objects.
# only py-amqp supports ssl and doesn't recognize amqps
# so fix up the connection string accordingly
connString = 'amqp://{0}:{1}@{2}:{3}/{4}'.format(options.mquser, options.mqpassword, options.mqserver, options.mqport, options.mqvhost)
if options.mqprotocol == 'amqps':
mqSSL = True
else:
mqSSL = False
mqConn = Connection(connString, ssl=mqSSL)
# Task Exchange for events sent via http for us to normalize and post to elastic search
if options.mqack:
# conservative, store msgs to disk, ack each message
eventTaskExchange = Exchange(name=options.taskexchange, type='direct', durable=True, delivery_mode=2)
else:
# fast, transient delivery, store in memory only, auto-ack messages
eventTaskExchange = Exchange(name=options.taskexchange, type='direct', durable=True, delivery_mode=1)
eventTaskExchange(mqConn).declare()
# Queue for the exchange
if options.mqack:
eventTaskQueue = Queue(options.taskexchange, exchange=eventTaskExchange, routing_key=options.taskexchange, durable=True, no_ack=False)
else:
eventTaskQueue = Queue(options.taskexchange, exchange=eventTaskExchange, routing_key=options.taskexchange, durable=True, no_ack=True)
eventTaskQueue(mqConn).declare()
# topic exchange for anyone who wants to queue and listen for mozdef.event
eventTopicExchange = Exchange(name=options.eventexchange, type='topic', durable=False, delivery_mode=1)
eventTopicExchange(mqConn).declare()
if hasUWSGI:
logger.info("started as uwsgi mule {0}".format(uwsgi.mule_id()))
else:
logger.info('started without uwsgi')
# consume our queue and publish on the topic exchange
taskConsumer(mqConn, eventTaskQueue, eventTopicExchange, es).run()
def flagdata(self, flag_file, reasons="any"):
logger.info(Color.HEADER + "Flagging " + reasons + " reasons" +
Color.ENDC)
script_path = 'casa_scripts/flag.py'
script_parameters = "{0} {1} {2}".format(self._dataset_path, flag_file,
reasons)
self._run(script_path, script_parameters)
def analyse_antennas_on_angular_dispersion(self):
logger.info(
Color.HEADER +
"Identifying bad Antennas based on angular dispersion in phases..."
+ Color.ENDC)
r_analyser = AngularDispersion(self.measurement_set, self)
r_analyser.identify_antennas_status()
def run(self):
self.taskQueue.set_message_class(RawMessage)
while True:
try:
records = self.taskQueue.get_messages(self.options.prefetch)
for msg in records:
msg_body = msg.get_body()
try:
# get_body() should be json
message_json = json.loads(msg_body)
self.on_message(message_json)
# delete message from queue
self.taskQueue.delete_message(msg)
except ValueError:
logger.error(
'Invalid message, not JSON <dropping message and continuing>: %r'
% msg_body)
self.taskQueue.delete_message(msg)
continue
time.sleep(.1)
except (SSLEOFError, SSLError, socket.error):
logger.info('Received network related error...reconnecting')
time.sleep(5)
self.connection, self.taskQueue = connect_sqs(
options.region, options.accesskey, options.secretkey,
options.taskexchange)
self.taskQueue.set_message_class(RawMessage)
def send_message_to_plugin(self, plugin_class, message, metadata=None):
if 'utctimestamp' in message and 'summary' in message:
message_log_str = u'{0} received message: ({1}) {2}'.format(
plugin_class.__module__, message['utctimestamp'],
message['summary'])
logger.info(message_log_str)
return plugin_class.onMessage(message), metadata
def phase_calibration(self):
logger.info(Color.SOURCE_HEADING + "Phase Calibration" + Color.ENDC)
phase_calibrator = PhaseCalibrator(self._measurement_set)
phase_calibrator.calibrate()
if config.CALIBRATION_STAGES['phase_calibration']['flagging']:
phase_calibrator.reduce_data()
else:
phase_calibrator.calibrate()
def flux_calibration(self):
logger.info(Color.SOURCE_HEADING + "Flux Calibration" + Color.ENDC)
flux_calibrator = FluxCalibrator(self._measurement_set)
flux_calibrator.run_setjy()
if config.CALIBRATION_STAGES['flux_calibration']['flagging']:
flux_calibrator.reduce_data()
else:
flux_calibrator.calibrate()
def stage_func_wrapper(*args):
if toggle_status:
start_time = datetime.datetime.now()
stage_func(*args)
end_time = datetime.datetime.now()
logger.info(Color.LightCyan + Color.UNDERLINE +
'Time spent in ' + stage_func.__name__ + '= ' +
str(abs((end_time - start_time).seconds)) +
" seconds" + Color.ENDC)
def base_image(self):
logger.info(Color.HEADER +
"Creating base image for {0}".format(self._dataset_path) +
Color.ENDC)
script_path = 'casa_scripts/base_image.py'
script_parameters = "{0} {1} {2}".format(self._dataset_path,
self._output_path,
config.CONFIG_PATH)
self._run(script_path, script_parameters)
def bandpass_calibration(self):
logger.info(Color.SOURCE_HEADING + "Bandpass Calibration" + Color.ENDC)
bandpass_calibrator = BandpassCalibrator(self._measurement_set)
bandpass_calibrator.calibrate()
if config.CALIBRATION_STAGES['bandpass_calibration'][
'run_auto_flagging']:
bandpass_calibrator.run_tfcrop()
bandpass_calibrator.run_rflag()
bandpass_calibrator.calibrate()
def r_flag(self, source_type, source_ids):
script_path = 'casa_scripts/r_flag.py'
source_ids = ','.join([str(source_id) for source_id in source_ids])
script_parameters = "{0} {1} {2} {3} {4}".format(
self._dataset_path, source_type, source_ids, config.CONFIG_PATH,
config.GLOBAL_CONFIGS['spw_range'])
logger.info(Color.HEADER + "Running Rflag auto-flagging algorithm" +
Color.ENDC)
self._run(script_path, script_parameters)
def train(self):
with tf.Session() as session:
logger.info('Initializing {} model...'.format(self._model_to_use))
model = self._models[self._model_to_use]()
logger.info('Preparing model...')
input = tf.placeholder(tf.float32, shape=[None, 784], name='input')
train = tf.placeholder(tf.bool)
labels = tf.placeholder(tf.float32, shape=[None, 10])
inference = model.infer(input, train)
cost = model.cost(inference, labels)
optimize = model.optimize(cost, 0.1)
evaluate = model.evaluate(inference, labels)
logger.info('Initializing variables...')
session.run(tf.global_variables_initializer())
logger.info('Initializing saver...')
self._saver = tf.train.Saver()
tf.summary.scalar('cost', cost)
writer = tf.summary.FileWriter(self._tensorboard_log_dir,
graph=tf.get_default_graph())
summarize = tf.summary.merge_all()
total_batch = int(self._mnist_data.train.num_examples /
self._batch_size)
step = 0
for epoch in range(self._epochs):
for _ in range(total_batch):
batch_x, batch_y = self._mnist_data.train.next_batch(
self._batch_size)
_, cost_value, summary = session.run(
[optimize, cost, summarize],
feed_dict={
input: batch_x,
labels: batch_y,
train: True,
})
step += 1
if step % 100 == 0:
logger.debug(
'epoch: {}, step: {}, cost: {:.2f}'.format(
epoch, step, cost_value))
writer.add_summary(summary, step)
accuracy = session.run(evaluate,
feed_dict={
input: self._mnist_data.test.images,
labels: self._mnist_data.test.labels,
train: False,
})
logger.info('The accuracy is: {:.2f}%'.format(accuracy * 100))
self._save_model(session)
def fourier_transform(self, field_name, cal_mode, spw_range, loop_count):
logger.info(Color.HEADER +
"Calculating fourier transform on {0}".format(field_name) +
Color.ENDC)
script_path = 'casa_scripts/fourier_transform.py'
script_parameters = "{0} {1} {2} {3} {4} {5}".format(
spw_range, self._output_path, cal_mode, loop_count,
self._dataset_path, field_name)
self._run(script_path, script_parameters)
def run(self):
self.taskQueue.set_message_class(RawMessage)
while True:
try:
records = self.taskQueue.get_messages(options.prefetch)
for msg in records:
body_message = msg.get_body()
event = json.loads(body_message)
if not event['Message']:
logger.error(
'Invalid message format for cloudtrail SQS messages'
)
logger.error('Malformed Message: %r' % body_message)
continue
if event['Message'] == 'CloudTrail validation message.':
# We don't care about these messages
continue
message_json = json.loads(event['Message'])
if 's3ObjectKey' not in message_json.keys():
logger.error(
'Invalid message format, expecting an s3ObjectKey in Message'
)
logger.error('Malformed Message: %r' % body_message)
continue
s3_log_files = message_json['s3ObjectKey']
for log_file in s3_log_files:
logger.debug('Downloading and parsing ' + log_file)
bucket = self.s3_connection.get_bucket(
message_json['s3Bucket'])
log_file_lookup = bucket.lookup(log_file)
events = self.process_file(log_file_lookup)
for event in events:
self.on_message(event)
self.taskQueue.delete_message(msg)
except KeyboardInterrupt:
sys.exit(1)
except Exception as e:
logger.exception(e)
time.sleep(3)
except (SSLEOFError, SSLError, socket.error) as e:
logger.info('Received network related error...reconnecting')
time.sleep(5)
self.connection, self.taskQueue = connect_sqs(
options.region, options.accesskey, options.secretkey,
options.taskexchange)
self.taskQueue.set_message_class(RawMessage)
def setjy(self, source_id, source_name):
logger.info(Color.HEADER + 'Running setjy' + Color.ENDC)
script_path = 'casa_scripts/setjy.py'
freq_band = "L"
model_path = "{0}/{1}_{2}.im".format(
config.CASA_CONFIGS['casa'][platform.system()]['model_path'],
source_name.split("_")[0], freq_band)
script_parameters = "{0} {1} {2} {3}".format(
config.GLOBAL_CONFIGS['spw_range'], self._dataset_path, source_id,
model_path)
self._run(script_path, script_parameters)
def main():
if hasUWSGI:
logger.info("started as uwsgi mule {0}".format(uwsgi.mule_id()))
else:
logger.info('started without uwsgi')
# establish api interface with papertrail
ptRequestor = PTRequestor(options.ptapikey, evmax=options.ptquerymax)
# consume our queue
taskConsumer(ptRequestor, es).run()
def create_line_image(self, calmode_config, parent_source_id):
logger.info(Color.HEADER +
"Creating line image at {0}".format(self._output_path) +
Color.ENDC)
script_path = 'casa_scripts/create_line_image.py'
cont_mode = 'ref'
continuum_image_model = self._last_continuum_image_model(
calmode_config, parent_source_id, cont_mode)
script_parameters = "{0} {1} {2} {3} {4}".format(
config.GLOBAL_CONFIGS['spw_range'], self._dataset_path,
self._output_path, continuum_image_model, config.CONFIG_PATH)
self._run(script_path, script_parameters)
def flag_and_calibrate_in_detail(self):
logger.info(Color.HEADER + "Started Detail Flagging..." + Color.ENDC)
detailed_analyser = DetailedAnalyser(self.measurement_set, self.config,
self.flag_file)
self._flag_bad_time(BAD_TIME, detailed_analyser.analyse_time, True)
self._flag_bad_time(BAD_ANTENNA_TIME,
detailed_analyser.analyse_antennas, False)
self._flag_bad_time(BAD_BASELINE_TIME,
detailed_analyser.analyse_baselines, False)
scan_ids = self.measurement_set.scan_ids(self.source_ids)
self.measurement_set.casa_runner.generate_flag_summary(
"detailed_flagging", scan_ids, self.source_type)
def generate_report(self, scan_ids):
logger.info("AntennaId, Polarisation, ScanId, R_Status, CP_Status")
for ant in self.__antennas:
for state in ant.get_states():
if state.scan_id in scan_ids and (
state.get_R_phase_status() == AntennaStatus.BAD and
state.get_closure_phase_status() == AntennaStatus.BAD):
logger.info(
" {0}\t \t{1}\t {2}\t {3}\t {4}".format(
ant.id, state.polarization, state.scan_id,
state.get_R_phase_status(),
state.get_closure_phase_status()))
def apply_target_source_calibration(self, source_id):
logger.info(Color.HEADER + "Applying Calibration to Target Source..." +
Color.ENDC)
flux_cal_fields = ",".join(
map(str, config.GLOBAL_CONFIGS['flux_cal_fields']))
phase_cal_fields = ",".join(
map(str, config.GLOBAL_CONFIGS['target_phase_src_map'][source_id]))
script_path = 'casa_scripts/target_source_calibration.py'
script_parameters = "{0} {1} {2} {3} {4}".format(
self._dataset_path, self._output_path, flux_cal_fields,
phase_cal_fields, source_id)
self._run(script_path, script_parameters)
def _create_all_spw_continuum_image(self, line_source):
if not self._is_single_spw_present():
cont_mode = 'spw'
spw_range = config.GLOBAL_CONFIGS['spw_range']
continuum_source = line_source.continuum(spw_range, cont_mode)
if config.TARGET_SOURCE_STAGES['all_spw']['continuum']['selfcal']:
continuum_source.base_image()
continuum_source.self_calibrate(cont_mode)
else:
logger.info(
Color.HEADER +
'Spw continuum image is already created [spw range contains only one spw]'
+ Color.ENDC)
def main(dataset_path):
start_time = datetime.datetime.now()
measurement_set = MeasurementSet(dataset_path, config.OUTPUT_PATH)
pipeline_stages = PipelineStage(measurement_set)
pipeline_stages.flag_known_bad_antennas()
pipeline_stages.flux_calibration()
pipeline_stages.bandpass_calibration()
pipeline_stages.phase_calibration()
pipeline_stages.target_source()
end_time = datetime.datetime.now()
logger.info(Color.UNDERLINE + 'Total time =' + str(abs((end_time - start_time).seconds)) + " seconds" + Color.ENDC)
def apply_line_calibration(self, calmode_config, parent_source_id, mode):
logger.info(Color.HEADER + "Applying calibration on Line.." +
Color.ENDC)
script_path = 'casa_scripts/apply_line_calibration.py'
p_loop_count = calmode_config["p"]["loop_count"]
ap_loop_count = calmode_config["ap"]["loop_count"]
p_table = '{0}/self_caled_p_{1}_{2}/p_selfcaltable_{3}.gcal'.format(
config.OUTPUT_PATH, mode, parent_source_id, p_loop_count)
ap_table = '{0}/self_caled_ap_{1}_{2}/ap_selfcaltable_{3}.gcal'.format(
config.OUTPUT_PATH, mode, parent_source_id, ap_loop_count)
script_parameters = "{0} {1} {2} {3} {4}".format(
p_loop_count, ap_loop_count, ap_table, p_table, self._dataset_path)
self._run(script_path, script_parameters)
def target_source(self):
if self._target_source_toggle():
for source_id in config.GLOBAL_CONFIGS['target_src_fields']:
logger.info(Color.SOURCE_HEADING +
"Target Source Calibration" + Color.ENDC)
target_source = TargetSource(self._measurement_set, source_id)
if config.TARGET_SOURCE_STAGES['calibrate']:
target_source.calibrate()
line = target_source.line()
line.measurement_set.generate_flag_summary("known_flags")
self._create_ref_continuum_image(line)
self._run_autoflagging_on_line(line)
self._create_all_spw_continuum_image(line)
self._create_line_image(line)
def apply_phase_calibration(self, flux_cal_field, source_config):
calib_params = CalibParams(*source_config['calib_params'])
logger.info(Color.HEADER + "Applying Phase Calibration..." +
Color.ENDC)
script_path = 'casa_scripts/phase_calibration.py'
phase_cal_fields = ",".join(
map(str, config.GLOBAL_CONFIGS['phase_cal_fields']))
refant = config.GLOBAL_CONFIGS['refant']
spw = format_spw_with_channels(config.GLOBAL_CONFIGS['spw_range'],
source_config['channels_to_avg'])
script_parameters = "{0} {1} {2} {3} {4} {5} {6} {7}".format(
self._dataset_path, self._output_path, flux_cal_field,
phase_cal_fields, spw, refant, calib_params.minsnr,
calib_params.solint)
self._run(script_path, script_parameters)
def identify_plugins(self, enabled_plugins):
if not os.path.exists(self.plugin_location):
return []
module_name = os.path.basename(self.plugin_location)
root_plugin_directory = os.path.join(self.plugin_location, '..')
plugin_manager = pynsive.PluginManager()
plugin_manager.plug_into(root_plugin_directory)
plugins = []
found_modules = pynsive.list_modules(module_name)
for found_module in found_modules:
module_filename, module_name = found_module.split('.')
if enabled_plugins is not None and module_name not in enabled_plugins:
# Skip this plugin since it's not listed as enabled plugins
# as long as we have specified some enabled plugins though
# this allows us to specify no specific plugins and get all of them
continue
try:
module_obj = pynsive.import_module(found_module)
reload(module_obj)
plugin_class_obj = module_obj.message()
if 'priority' in dir(plugin_class_obj):
priority = plugin_class_obj.priority
else:
priority = 100
logger.info('[*] plugin {0} registered to receive messages with {1}'.format(module_name, plugin_class_obj.registration))
plugins.append(
{
'plugin_class': plugin_class_obj,
'registration': plugin_class_obj.registration,
'priority': priority
}
)
except Exception as e:
logger.exception('Received exception when loading {0} plugins\n{1}'.format(module_name, e.message))
plugin_manager.destroy()
return plugins
