added set_num_threads (1a03e743) · Commits · Bo Zhang / csst

csst/msc/inst_corr.py

+7 −2

Original line number	Original line	Diff line number	Diff line
	from pathlib import Path		from pathlib import Path

			import torch
	import numpy as np		import numpy as np
	from ccdproc import cosmicray_lacosmic		from ccdproc import cosmicray_lacosmic
	from deepCR import deepCR		from deepCR import deepCR
	@@ -24,6 +25,9 @@ class CsstMscInstrumentProc(CsstProcessor):
	self.__wht = None		self.__wht = None
	self.__flg = None		self.__flg = None

			def set_num_threads(self, n_threads=1):
			torch.set_num_threads(n_threads)

	def _do_fix(self, raw, bias, dark, flat):		def _do_fix(self, raw, bias, dark, flat):
	'''仪器效应改正		'''仪器效应改正

	@@ -100,7 +104,7 @@ class CsstMscInstrumentProc(CsstProcessor):
	n_jobs=self.n_jobs)		n_jobs=self.n_jobs)
	else:		else:
	masked, cleaned = model.clean(		masked, cleaned = model.clean(
	self.__img, threshold=0.5, inpaint=True, segment=False, patch=256, parallel=False,		self.__img, threshold=0.5, inpaint=True, segment=True, patch=256, parallel=False,
	n_jobs=self.n_jobs)		n_jobs=self.n_jobs)
	else:		else:
	cleaned, masked = cosmicray_lacosmic(ccd=self.__img,		cleaned, masked = cosmicray_lacosmic(ccd=self.__img,
	@@ -144,8 +148,9 @@ class CsstMscInstrumentProc(CsstProcessor):
	weight[self.__flg > 0] = 0		weight[self.__flg > 0] = 0
	self.__wht = weight		self.__wht = weight

	def prepare(self, n_jobs=2, **kwargs):		def prepare(self, n_jobs=2, n_threads=1, **kwargs):
	self.n_jobs = n_jobs		self.n_jobs = n_jobs
			self.set_num_threads(n_threads)
	for name in kwargs:		for name in kwargs:
	self._switches[name] = kwargs[name]		self._switches[name] = kwargs[name]