fix cmg bug; improve the program run speed (95624292) · Commits · csst-planning / csst_survey_sim

survey_sim/constraints/_utils.py

+63 −1

Original line number	Diff line number	Diff line
		@@ -297,6 +297,24 @@ def calculate_trace(matrix):
		# 返回迹的总和
		return trace_sum

		# @njit(fastmath=True)
		def quaternion_from_axis_angle_batch(axes, angles):
		"""批量计算轴角到四元数"""
		angles_half = angles / 2
		sin_half = np.sin(angles_half)
		cos_half = np.cos(angles_half)
		return np.column_stack([
		cos_half,
		axes[:, 0] * sin_half,
		axes[:, 1] * sin_half,
		axes[:, 2] * sin_half
		])
		# return np.vstack((
		# cos_half,
		# axes[:, 0] * sin_half,
		# axes[:, 1] * sin_half,
		# axes[:, 2] * sin_half
		# )).T

		def quaternion_from_axis_angle(axis, angle_rad):
		"""
		@@ -309,7 +327,23 @@ def quaternion_from_axis_angle(axis, angle_rad):
		q_xyz = axis * np.sin(half_angle)
		return np.concatenate(([q0], q_xyz))


		# @njit(fastmath=True)
		def quaternion_multiply_batch(q1, q2):
		"""批量四元数乘法"""
		return np.column_stack([
		q1[:, 0] * q2[:, 0] - q1[:, 1] * q2[:, 1] - q1[:, 2] * q2[:, 2] - q1[:, 3] * q2[:, 3],
		q1[:, 0] * q2[:, 1] + q1[:, 1] * q2[:, 0] + q1[:, 2] * q2[:, 3] - q1[:, 3] * q2[:, 2],
		q1[:, 0] * q2[:, 2] - q1[:, 1] * q2[:, 3] + q1[:, 2] * q2[:, 0] + q1[:, 3] * q2[:, 1],
		q1[:, 0] * q2[:, 3] + q1[:, 1] * q2[:, 2] - q1[:, 2] * q2[:, 1] + q1[:, 3] * q2[:, 0]
		])
		# return np.vstack((
		# q1[:, 0] * q2[:, 0] - q1[:, 1] * q2[:, 1] - q1[:, 2] * q2[:, 2] - q1[:, 3] * q2[:, 3],
		# q1[:, 0] * q2[:, 1] + q1[:, 1] * q2[:, 0] + q1[:, 2] * q2[:, 3] - q1[:, 3] * q2[:, 2],
		# q1[:, 0] * q2[:, 2] - q1[:, 1] * q2[:, 3] + q1[:, 2] * q2[:, 0] + q1[:, 3] * q2[:, 1],
		# q1[:, 0] * q2[:, 3] + q1[:, 1] * q2[:, 2] - q1[:, 2] * q2[:, 1] + q1[:, 3] * q2[:, 0]
		# )).T

		# @njit(fastmath=True)
		def quaternion_multiply(q1, q2):
		"""
		计算两个四元数的乘积 q = q1 * q2
		@@ -458,6 +492,34 @@ def transfer2lonlat_quaternion_noPa(lon=90., lat=90.):

		return q_total

		def rotate_quaternion_byself_batch(lon=90., lat=90., pa=0., quaternion=[1, 0, 0, 0]):
		"""支持标量或向量输入，返回旋转后的四元数"""
		if isinstance(lon, (float, int)):
		# 标量模式（兼容旧代码）
		lon_rad = radians(lon)
		lat_rad = radians(lat)
		pa_rad = radians(pa)
		axis = np.array([
		cos(lat_rad) * cos(lon_rad),
		cos(lat_rad) * sin(lon_rad),
		sin(lat_rad)
		])
		q_t = quaternion_from_axis_angle(axis, pa_rad)
		return quaternion_multiply(q_t, quaternion)
		else:
		# 向量模式（批量计算）
		lon_rad = np.radians(lon)
		lat_rad = np.radians(lat)
		pa_rad = np.radians(pa)

		axis = np.column_stack([
		np.cos(lat_rad) * np.cos(lon_rad),
		np.cos(lat_rad) * np.sin(lon_rad),
		np.sin(lat_rad)
		])

		q_t = quaternion_from_axis_angle_batch(axis, pa_rad)
		return quaternion_multiply_batch(q_t, quaternion)

		def rotate_quaternion_byself(lon=90., lat=90., pa=0., quaternion=[1, 0, 0, 0]):
		lon_rad = math.radians(lon)

survey_sim/constraints/cmg_constraint.py

+58 −6

Original line number	Diff line number	Diff line
		@@ -52,6 +52,22 @@ class cmgConstraint(object):
		cmg_use = self.CMG_weight_single[-1]
		return cmg_use

		def get_cmg_use_array(self, tAngles):
		cmg_use = np.zeros(len(tAngles))
		tAngles[tAngles<0] = 0
		tAngles[tAngles>180] = 180

		for i in np.arange(len(self.CMG_angel_range)):
		c_range = self.CMG_angel_range[i]
		w = self.CMG_weight_single[i]

		ids_t = (tAngles >= c_range[0]) & (tAngles < c_range[1])
		cmg_use[ids_t] = w

		ids_t = tAngles == 180
		cmg_use[ids_t] = self.CMG_weight_single[-1]
		return cmg_use

		def initCMGWeight(self, flag=1):
		CMG_TM_FLAG = flag
		self.cmg_tm_tAngle = np.zeros(10)
		@@ -236,7 +252,7 @@ class cmgConstraint(object):

		def computeCMGTotal_newComein(
		self, curTime=2459766.0, tTime=100, exTime=150, cmg_use=0.5):
		cmg_times = 0
		cmg_times = (tTime + exTime) / 86400.0
		# 用于获取当前队列对应的时间长度

		if self.cmgNodes:
		@@ -251,9 +267,45 @@ class cmgConstraint(object):
		return cmg_total_use
		# 如果当前队列的时间长度超过了一轨的时间，就将队列头部的元素给pop掉，并重新计算CMG的总消耗量
		if self.cmgNodes:
		for cNode in self.cmgNodes:
		nodeSize = len(self.cmgNodes)
		if nodeSize > 1:
		for i in np.arange(1,nodeSize,1):
		cNode = self.cmgNodes[i]
		pNode = self.cmgNodes[i-1]
		cmg_times = curTime - cNode[0] + (tTime + exTime)/ 86400.0
		cmg_total_use = cmg_total_use - cNode[1]
		cmg_total_use = cmg_total_use - pNode[1]
		if cmg_times <= self.orbit_time:
		break
		return cmg_total_use

		def computeCMGTotal_newComein_array(
		self, curTime=2459766.0, tTimes=[100], exTimes=[150], cmg_uses=[0.5]):
		# cmg_times = np.zeros(len(tTimes))
		cmg_times = (tTimes + exTimes) / 86400.0
		# 用于获取当前队列对应的时间长度

		if self.cmgNodes:
		cmg_times = curTime - \
		self.cmgNodes[0][0] + (tTimes + exTimes) / 86400.0
		# ????? 这里是否应该包含exTime ?????
		# else:

		cmg_total_use = self.total_cmg + cmg_uses

		ids = cmg_times > self.orbit_time
		if(np.sum(ids)>0):
		if self.cmgNodes:
		nodeSize = len(self.cmgNodes)
		if nodeSize > 1:
		for i in np.arange(1,nodeSize,1):
		cNode = self.cmgNodes[i]
		pNode = self.cmgNodes[i-1]
		cmg_times[ids] = curTime - cNode[0] + (tTimes[ids] + exTimes[ids]) / 86400.0
		cmg_total_use[ids] = cmg_total_use[ids] - pNode[1]
		ids = cmg_times > self.orbit_time
		if np.sum(ids)==0:
		break
		if(np.sum(ids)>0):
		cmg_times[ids] = (tTimes[ids] + exTimes[ids]) / 86400.0
		cmg_total_use[ids] = cmg_total_use[ids] - self.cmgNodes[-1][1]
		return cmg_total_use

survey_sim/constraints/obscure_constraint.py

+149 −1

Original line number	Diff line number	Diff line
		@@ -16,6 +16,8 @@ from survey_sim.satOrbit import locateSat, loadSatOrbitDat

		from survey_sim.constraints import _utils

		from numba import njit

		MAX_VAL = 100000

		"""
		@@ -226,6 +228,83 @@ def IsObscureByEarth(sat=None, sun=None, p=None, constr=None):
		return isObscure, angleValue


		def IsObscureByEarth_array(sat=None, sun=None, p=None, constr=None):
		"""
		判断卫星观测点是否被地球遮挡（支持多个卫星和观测点）

		Args:
		sat (np.array): (n, 3) 卫星位置向量（ECI坐标系）
		sun (np.array): (3,) 太阳位置向量（ECI坐标系）
		p (np.array): (n, 3) 观测点向量（ECI坐标系）
		constr (object): 包含光照条件的阈值参数

		Returns:
		np.array: (n,) 返回每个观测点的遮挡情况（1=遮挡，0=可见）
		"""
		if sat.ndim == 1:
		sat = sat.reshape(1, 3)
		if p.ndim == 1:
		p = p.reshape(1, 3)

		n = sat.shape[0] # 卫星/观测点的数量

		# 计算向量模长
		modSat = np.linalg.norm(sat, axis=1) # (n,)
		modPoint = np.linalg.norm(p, axis=1) # (n,)
		modSun = np.linalg.norm(sun) # 标量

		# 计算卫星天顶角余弦 (n,)
		withLocalZenithAngle = np.sum(p * sat, axis=1) / (modPoint * modSat)

		# 计算太阳与卫星的夹角余弦 (n,)
		innerM_sat_sun = np.sum(sat * sun, axis=1) # (n,)
		cosAngle = innerM_sat_sun / (modSat * modSun) # (n,)

		# 定义地球阴影边界条件
		tanPE = 0.3598250103051328 # tan(19.79°)
		cosPE1 = -0.3385737 # cos(109.79°)
		cosPE2 = 0.3385737 # cos(70.21°)

		# 判断每个卫星是否在阳照区、阴影区或过渡区 (n,)
		isInSunSide = np.zeros(n, dtype=int)
		isInSunSide[cosAngle < cosPE1] = -1 # 完全阴影区
		isInSunSide[cosAngle > cosPE2] = 1 # 完全阳照区
		# 其余情况为过渡区（isInSunSide=0）

		# 计算观测点与太阳的夹角余弦 (n,)
		p_to_sun_cos = np.sum(p * sun, axis=1) / (modPoint * modSun)

		# 初始化遮挡结果 (n,)
		isObscure = np.ones(n, dtype=int)

		# 完全阳照区
		mask_light = (isInSunSide == 1)
		isObscure[mask_light & (withLocalZenithAngle >= constr.satZenith_angle_light_max_cos)] = 0

		# 完全阴影区
		mask_dark = (isInSunSide == -1)
		isObscure[mask_dark & (withLocalZenithAngle >= constr.satZenith_angle_dark_cos)] = 0

		# 过渡区
		mask_transition = (isInSunSide == 0)
		# 观测点在太阳方向（p_to_sun_cos >= 0）
		mask_sun_side = (p_to_sun_cos >= 0)
		isObscure[mask_transition & mask_sun_side & (withLocalZenithAngle >= constr.satZenith_angle_light_max_cos)] = 0
		# 观测点在地球阴影方向（p_to_sun_cos < 0）
		isObscure[mask_transition & ~mask_sun_side & (withLocalZenithAngle >= constr.satZenith_angle_dark_cos)] = 0

		return isObscure

		@njit(fastmath=True)
		def is_in_sun_side_numba(sat, sun):
		"""Numba 加速的阳照区判断"""
		inner_product = sat[0] * sun[0] + sat[1] * sun[1] + sat[2] * sun[2]
		mod_sat = sat[0]2 + sat[1]2 + sat[2]**2
		mod_sun = sun[0]2 + sun[1]2 + sun[2]**2
		cos_angle = inner_product / np.sqrt(mod_sat * mod_sun)
		return int(cos_angle >= -0.3385737) # cos(109.79°)


		# /*
		# * 根据卫星和太阳的位置来判断是否在阳照区，在阳照区为1，否则为0
		# */
		@@ -242,6 +321,48 @@ def IsInSunSide(sat=None, sun=None):
		in_sun_side = (int)(cosAngle >= -0.3385737)
		return in_sun_side

		def aquire_shadow_time_fast(curTime, orbitDat, ephlib, max_steps=1000, tol=1e-6):
		"""二分法快速定位阴影区边界"""
		# Step 1: 初始粗略搜索边界
		t_start = curTime
		t_end = curTime + 0.08333333333333333 # 假设最大搜索范围2小时，120分钟

		# 检查初始状态
		sat_start, _, _ = locateSat(t_start, orbitDat)
		sun_start = locate_sun(t_start, ephlib)
		initial_state = is_in_sun_side_numba(sat_start, sun_start)

		# Step 2: 二分法精确定位
		if initial_state == 0:
		# 当前在阴影区，找结束时间
		left, right = t_start, t_end
		for _ in range(max_steps):
		mid = (left + right) / 2
		sat_mid, _, _ = locateSat(mid, orbitDat)
		sun_mid = locate_sun(mid, ephlib)
		state = is_in_sun_side_numba(sat_mid, sun_mid)
		if state == 0:
		left = mid
		else:
		right = mid
		if right - left < tol:
		return [t_start, right]
		else:
		# 当前在阳照区，找开始时间
		left, right = t_start, t_end
		for _ in range(max_steps):
		mid = (left + right) / 2
		sat_mid, _, _ = locateSat(mid, orbitDat)
		sun_mid = locate_sun(mid, ephlib)
		state = is_in_sun_side_numba(sat_mid, sun_mid)
		if state == 1:
		left = mid
		else:
		right = mid
		if right - left < tol:
		return [right, t_end]

		return [t_start, t_end] # 默认返回（实际应处理未收敛情况）

		# /*
		# * 根据进入和离开阴影区的时间来判断是否在阳照区，在阳照区为1，否则为0
		@@ -348,4 +469,31 @@ def IsObscureSolarPlane(skymap_ids=None, sun_ecl_car=[0., 0., 0.], dist_sun=1.5e

		# return skymap_ids[norm_ids1]
		# return np.stack((skymap_ids[norm_ids1], value_sun_angle_to_nomals[norm_ids1]), axis=1)
		return skymap_ids[norm_ids1], value_sun_angle_to_nomals[norm_ids1]
		return skymap_ids[norm_ids1], value_sun_angle_to_nomals[norm_ids1], norm_ids1

		def IsObscureSolarPlane_batch(skymap_ids=None, sun_ecl_car=[0., 0., 0.], dist_sun=1.5e8, skyMap=None, cos_sun_plane_angle=-1.0):
		# 1. 计算太阳角度
		satPlane_norms = np.array([skyMap.skymap[k].solar_plane_norm for k in skymap_ids]) # 转换为 NumPy 数组
		value_sun_angle_to_nomals = np.dot(satPlane_norms, sun_ecl_car) / dist_sun
		norm_ids = value_sun_angle_to_nomals < 0
		value_sun_angle_to_nomals[norm_ids] = -value_sun_angle_to_nomals[norm_ids]
		norm_ids1 = value_sun_angle_to_nomals >= cos_sun_plane_angle

		# 2. 批量更新四元数（避免循环）
		mask = norm_ids & norm_ids1
		if np.any(mask):
		# 提取所有需要旋转的观测单元
		obseveSkyUnits = [skyMap.skymap[id] for id in skymap_ids[mask]]
		lons = np.array([unit.ecl_lon for unit in obseveSkyUnits])
		lats = np.array([unit.ecl_lat for unit in obseveSkyUnits])
		pas = np.full_like(lons, 180.) # pa=180
		quaternions = np.array([unit.p_quaternion for unit in obseveSkyUnits])

		# 批量旋转（假设 rotate_quaternion_byself 支持向量化）
		rotated_quaternions = _utils.rotate_quaternion_byself_batch(lon=lons, lat=lats, pa=pas, quaternion=quaternions)

		# 更新 quaternion
		for i, unit in enumerate(obseveSkyUnits):
		unit.p_quaternion = rotated_quaternions[i]

		return skymap_ids[norm_ids1], value_sun_angle_to_nomals[norm_ids1], norm_ids1
		No newline at end of file

survey_sim/ephemeris/jplFileReader/ephJPLDE405

0 → 100755

+61.6 KiB

File added.

No diff preview for this file type.

View file

survey_sim/strategies/SurveySim_MSC_mpi.py

+226 −351

File changed.

Preview size limit exceeded, changes collapsed.