[Code scan] Honor TensorFlow change-bias --numb-batch

[njzjz]

This issue comes from a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.

Problem

change-bias --numb-batch is documented as the number of frames to use per data system, with 0 meaning all data:

deepmd-kit/deepmd/main.py

Lines 864 to 870 in 73de44b

	parser_change_bias.add_argument(
	"-n",
	"--numb-batch",
	default=0,
	type=int,
	help="The number of frames for bias changing in one data system. 0 means all data.",
	)

The TensorFlow change-bias call stack threads numb_batch through _change_bias_checkpoint_file():

deepmd-kit/deepmd/tf/entrypoints/change_bias.py

Lines 134 to 144 in 73de44b

	def _change_bias_checkpoint_file(
	checkpoint_prefix: str,
	mode: str,
	bias_value: list | None,
	datafile: str | None,
	system: str,
	numb_batch: int,
	model_branch: str | None,
	output: str | None,
	log_level: int,
	) -> None:

But the data-based bias path drops it and calls _apply_data_based_bias(...) without numb_batch:

deepmd-kit/deepmd/tf/entrypoints/change_bias.py

Lines 231 to 245 in 73de44b

	if bias_value is not None:
	# Use user-defined bias
	_apply_user_defined_bias(trainer, bias_value)
	else:
	# Use data-based bias calculation
	type_map = data.get_type_map()
	if len(type_map) == 0:
	# If data doesn't have type_map, get from model
	type_map = trainer.model.get_type_map()
	log.info(f"Changing bias for model with type_map: {type_map}")
	log.info(f"Using bias adjustment mode: {bias_adjust_mode}")
	# Read current bias values from the session (after variables are restored)
	_apply_data_based_bias(trainer, data, type_map, bias_adjust_mode)

_apply_data_based_bias() then hard-codes ntest=1 when computing the new bias:

deepmd-kit/deepmd/tf/entrypoints/change_bias.py

Lines 377 to 397 in 73de44b

	# Create a temporary frozen model to use with change_energy_bias_lower
	with tempfile.NamedTemporaryFile(suffix=".pb", delete=False) as temp_frozen:
	freeze(
	checkpoint_folder=str(Path(trainer.run_opt.init_model).parent),
	output=temp_frozen.name,
	)
	try:
	# Create DeepPotential object for evaluation
	dp = DeepPotential(temp_frozen.name)
	# Use change_energy_bias_lower with the current bias values from session
	new_bias = change_energy_bias_lower(
	data,
	dp,
	type_map, # origin_type_map
	type_map, # full_type_map
	current_bias, # Use the restored bias values
	bias_adjust_mode=bias_adjust_mode,
	ntest=1,
	)

and change_energy_bias_lower() caps each system to min(nframes, ntest):

deepmd-kit/deepmd/tf/fit/ener.py

Lines 1040 to 1044 in 73de44b

	for sys in data.data_systems:
	test_data = sys.get_test()
	nframes = test_data["box"].shape[0]
	numb_test = min(nframes, ntest)
	if mixed_type:

Impact

TensorFlow data-based change-bias always uses one frame per system regardless of -n/--numb-batch. Passing -n 10 or relying on the documented 0 means all data behavior does not affect the bias estimate.

Suggested fix

Pass numb_batch into _apply_data_based_bias() and forward it to change_energy_bias_lower() as ntest, preserving the existing 0 means all data behavior if that helper needs a sentinel translation.