113class MazeDataset(GPTDataset[MazeDatasetConfig]):
114	"""a maze dataset class. This is a collection of solved mazes, and should be initialized via `MazeDataset.from_config`"""
115
116	def __init__(
117		self,
118		cfg: MazeDatasetConfig,
119		mazes: typing.Sequence[SolvedMaze],
120		generation_metadata_collected: dict | None = None,
121	) -> None:
122		"""initialize a maze dataset from a config and a list of solved mazes"""
123		super().__init__()
124		self.cfg: MazeDatasetConfig = cfg
125		self.mazes: list[SolvedMaze] = list(mazes)
126		self.generation_metadata_collected: dict | None = generation_metadata_collected
127
128	# TYPING: error: Return type "MazeDataset" of "from_config" incompatible with return type "T_Dataset" in supertype "GPTDataset"  [override]
129	@classmethod
130	def from_config(  # type: ignore[override]
131		cls,
132		# TYPING: error: Argument 1 of "from_config" is incompatible with supertype "GPTDataset"; supertype defines the argument type as "T_DatasetConfig"  [override]
133		cfg: MazeDatasetConfig,  # type: ignore[override]
134		do_generate: bool = True,
135		load_local: bool = True,
136		save_local: bool = True,
137		zanj: ZANJ | None = None,
138		do_download: bool = True,
139		local_base_path: Path = Path("data/maze_dataset"),
140		except_on_config_mismatch: bool = True,
141		allow_generation_metadata_filter_mismatch: bool = True,
142		verbose: bool = False,
143		**kwargs,
144	) -> "MazeDataset":
145		"""create a maze dataset from a config
146
147		priority of loading:
148		1. load from local
149		2. download
150		3. generate
151
152		"""
153		return cast(
154			MazeDataset,
155			super().from_config(
156				cfg=cfg,
157				do_generate=do_generate,
158				load_local=load_local,
159				save_local=save_local,
160				zanj=zanj,
161				do_download=do_download,
162				local_base_path=local_base_path,
163				except_on_config_mismatch=except_on_config_mismatch,
164				allow_generation_metadata_filter_mismatch=allow_generation_metadata_filter_mismatch,
165				verbose=verbose,
166				**kwargs,
167			),
168		)
169
170	def data_hash(self) -> int:
171		"""return a hash of the data"""
172		return stable_hash(str(tuple([x.serialize() for x in self.mazes])))
173
174	def __getitem__(self, i: int) -> SolvedMaze:
175		"""get a maze by index"""
176		return self.mazes[i]
177
178	def __iter__(self) -> typing.Iterator[SolvedMaze]:
179		"""iterate over the mazes"""
180		return iter(self.mazes)
181
182	def __deepcopy__(self, memo) -> "MazeDataset":  # noqa: ANN001
183		"""deepcopy the dataset
184
185		FIX: this isnt actually a deepcopy I think?
186		"""
187		return MazeDataset.load(self._serialize_full())
188
189	# TYPING: get type hints on the tokenizer here
190	@overload
191	def as_tokens(
192		self,
193		maze_tokenizer,  # noqa: ANN001
194		limit: int | None = None,
195		join_tokens_individual_maze: Literal[False] = False,
196	) -> list[list[str]]: ...
197	@overload
198	def as_tokens(
199		self,
200		maze_tokenizer,  # noqa: ANN001
201		limit: int | None = None,
202		join_tokens_individual_maze: Literal[True] = True,
203	) -> list[str]: ...
204	def as_tokens(
205		self,
206		maze_tokenizer,  # TODO: MazeTokenizer
207		limit: int | None = None,
208		join_tokens_individual_maze: bool = False,
209	) -> list[list[str]] | list[str]:
210		"""return the dataset as tokens according to the passed `maze_tokenizer`
211
212		the `maze_tokenizer` should be either a `MazeTokenizer` or a `MazeTokenizerModular`
213
214		if `join_tokens_individual_maze` is True, then the tokens of each maze are
215		joined with a space, and the result is a list of strings.
216		i.e.:
217
218			>>> dataset.as_tokens(join_tokens_individual_maze=False)
219			[["a", "b", "c"], ["d", "e", "f"]]
220			>>> dataset.as_tokens(join_tokens_individual_maze=True)
221			["a b c", "d e f"]
222		"""
223		output: list[list[str]] = [
224			maze.as_tokens(maze_tokenizer) for maze in self.mazes[:limit]
225		]
226		if join_tokens_individual_maze:
227			return [" ".join(tokens) for tokens in output]
228		else:
229			return output
230
231	def __len__(self) -> int:
232		"""return the number of mazes in the dataset"""
233		return len(self.mazes)
234
235	def __eq__(self, other: object) -> bool:
236		"""compare two datasets"""
237		if not isinstance(other, MazeDataset):
238			raise NotImplementedError(
239				"can only compare with other MazeDataset objects",
240			)
241		# TODO: compare hashes of data instead of the data itself?
242		return self.cfg == other.cfg and self.mazes == other.mazes
243
244	def assert_equal(self, other: "MazeDataset") -> None:
245		"""assert that two datasets are equal"""
246		assert isinstance(other, MazeDataset)
247		assert self.cfg == other.cfg, f"{self.cfg.diff(other.cfg) = }"
248		assert self.mazes == other.mazes, f"{self.mazes = }, {other.mazes = }"
249
250	@classmethod
251	def generate(
252		cls,
253		cfg: MazeDatasetConfig,
254		gen_parallel: bool = False,
255		pool_kwargs: dict | None = None,
256		verbose: bool = False,
257		# TODO: what to do when unexpected kwargs are passed?
258		**kwargs,  # noqa: ARG003
259	) -> "MazeDataset":
260		"""Generate a maze dataset given a config and some generation parameters"""
261		# Copy the config to avoid modifying the original
262		cfg_cpy: MazeDatasetConfig = MazeDatasetConfig.load(
263			json.loads(json.dumps(cfg.serialize())),
264		)
265
266		if pool_kwargs is None:
267			pool_kwargs = dict()
268		maze_indexes: Int[np.ndarray, " maze_index"] = np.arange(cfg_cpy.n_mazes)  # type: ignore[assignment]
269
270		solved_mazes: list[SolvedMaze | None]
271		# Configure tqdm for progress bar
272		tqdm_kwargs: dict = dict(
273			total=cfg_cpy.n_mazes,
274			unit="maze",
275			desc="generating & solving mazes",
276			disable=not verbose,
277		)
278		# TODO: don't use the global unless generating in parallel!
279		if gen_parallel:
280			with multiprocessing.Pool(
281				**pool_kwargs,
282				initializer=_maze_gen_init_worker,
283				initargs=(cfg_cpy,),
284			) as pool:
285				solved_mazes = list(
286					tqdm.tqdm(
287						pool.imap(_generate_maze_helper, maze_indexes),
288						**tqdm_kwargs,
289					),
290				)
291
292		else:
293			_maze_gen_init_worker(cfg_cpy)
294			solved_mazes = list(
295				tqdm.tqdm(
296					map(
297						# TYPING:  error: Argument 1 to "map" has incompatible type "Callable[[int], SolvedMaze | None]"; expected "Callable[[str], SolvedMaze | None]"  [arg-type]
298						# why does it think tolist() returns a string?
299						_generate_maze_helper,  # type: ignore[arg-type]
300						maze_indexes.tolist(),
301					),
302					**tqdm_kwargs,
303				),
304			)
305
306		# Filter out None values explicitly after ensuring all results are collected
307		solved_mazes_: list[SolvedMaze] = [
308			maze for maze in solved_mazes if maze is not None
309		]
310		# solved_mazes_ = list(filter(lambda x: x is not None, solved_mazes))
311
312		# Update the config with the actual number of mazes
313		cfg_cpy.n_mazes = len(solved_mazes_)
314
315		dataset: MazeDataset = cls(
316			cfg=cfg_cpy,
317			mazes=solved_mazes_,
318		)
319
320		dataset.update_self_config()  # Call `update_self_config()` to ensure the dataset's config reflects changes
321
322		np.random.seed(cfg_cpy.seed)  # Reset the seed to the value in the config copy
323
324		return dataset
325
326	@classmethod
327	def download(cls, cfg: MazeDatasetConfig, **kwargs) -> "MazeDataset":
328		"(not implemented yet!) download a maze dataset from the internet"
329		raise NotImplementedError("not implemented yet")
330
331	@classmethod
332	def load(cls: "type[MazeDataset]", data: JSONdict) -> "MazeDataset":
333		"""load from zanj/json"""
334		if data[_FORMAT_KEY] == "MazeDataset:minimal":
335			return cls._load_minimal(data)
336		elif data[_FORMAT_KEY] == "MazeDataset:minimal_soln_cat":
337			return cls._load_minimal_soln_cat(data)
338		elif data[_FORMAT_KEY] == "MazeDataset":
339			if (
340				SERIALIZE_MINIMAL_THRESHOLD == -1
341			):  # Allow access to `_load_legacy` for profiling
342				return cls._load_legacy(data)
343			return cls._load_full(data)
344		else:
345			err_msg: str = f"`_FORMAT_KEY` string {data[_FORMAT_KEY] = } is not a recognized `MazeDataset` format. ({_FORMAT_KEY = })"
346			raise KeyError(
347				err_msg,
348			)
349
350	@classmethod
351	def _load_full(cls, data: JSONdict) -> "MazeDataset":
352		assert data[_FORMAT_KEY] == "MazeDataset"
353		return cls(
354			cfg=MazeDatasetConfig.load(data["cfg"]),  # type: ignore[arg-type]
355			mazes=load_item_recursive(data["mazes"], tuple()),
356			generation_metadata_collected=data["generation_metadata_collected"],  # type: ignore[arg-type]
357		)
358
359	@classmethod
360	def _load_minimal(cls, data: JSONdict) -> "MazeDataset":
361		assert data[_FORMAT_KEY] == "MazeDataset:minimal"
362		return cls(
363			cfg=MazeDatasetConfig.load(data["cfg"]),  # type: ignore[arg-type]
364			generation_metadata_collected=data["generation_metadata_collected"],  # type: ignore[arg-type]
365			mazes=[
366				SolvedMaze(
367					clist,
368					soln[:slen, ...],
369				)
370				for clist, slen, soln in zip(
371					load_item_recursive(data["maze_connection_lists"], tuple()),
372					load_item_recursive(data["maze_solution_lengths"], tuple()),
373					load_item_recursive(data["maze_solutions"], tuple()),
374					strict=False,
375					# load_item_recursive(data["maze_endpoints"], tuple()),
376				)
377			],
378		)
379
380	@classmethod
381	def _load_minimal_soln_cat(cls, data: JSONdict) -> "MazeDataset":
382		assert data[_FORMAT_KEY] == "MazeDataset:minimal_soln_cat"
383
384		maze_solution_lengths = load_item_recursive(
385			data["maze_solution_lengths"],
386			tuple(),
387		)
388		maze_solutions_concat = load_item_recursive(
389			data["maze_solutions_concat"],
390			tuple(),
391		)
392		maze_solutions = np.split(
393			maze_solutions_concat,
394			np.cumsum(maze_solution_lengths)[:-1],
395			axis=0,
396		)
397
398		return cls(
399			cfg=load_item_recursive(data["cfg"], tuple()),
400			generation_metadata_collected=load_item_recursive(
401				data["generation_metadata_collected"],
402				tuple(),
403			),
404			mazes=[
405				SolvedMaze(
406					connection_list=clist,
407					solution=soln,
408				)
409				for clist, soln in zip(
410					load_item_recursive(data["maze_connection_lists"], tuple()),
411					# load_item_recursive(data["maze_endpoints"], tuple()),
412					maze_solutions,
413					strict=False,
414				)
415			],
416		)
417
418	@classmethod
419	def _load_legacy(cls, data: JSONdict) -> "MazeDataset":
420		"""Legacy `load` method from <0.5.2. Used exclusively for profiling comparison."""
421		assert data[_FORMAT_KEY] == "MazeDataset"
422		return cls(
423			**{
424				key: load_item_recursive(data[key], tuple())
425				for key in ["cfg", "mazes", "generation_metadata_collected"]
426			},
427		)
428
429	def serialize(self) -> JSONdict:
430		"""serialize to zanj/json"""
431		if (
432			SERIALIZE_MINIMAL_THRESHOLD is not None
433			and len(self) >= SERIALIZE_MINIMAL_THRESHOLD
434		):
435			return self._serialize_minimal()
436		return self._serialize_full()
437
438	def _serialize_full(self) -> JSONdict:
439		return {
440			_FORMAT_KEY: "MazeDataset",
441			"cfg": json_serialize(self.cfg),
442			"fname": self.cfg.to_fname(),
443			"mazes": json_serialize(self.mazes),
444			"generation_metadata_collected": json_serialize(
445				self.generation_metadata_collected,
446			),
447		}
448
449	def _serialize_minimal(self) -> JSONdict:
450		"alternate serialization where metadata is collected and mazes are stored in concatenated form"
451		filtered_meta: MazeDataset
452		if self.generation_metadata_collected is None:
453			filtered_meta = self.filter_by.collect_generation_meta()
454		else:
455			filtered_meta = self
456
457		max_solution_len: int = max(m.solution.shape[0] for m in filtered_meta.mazes)
458		n_mazes: int = len(filtered_meta.mazes)
459		grid_n: int = filtered_meta.cfg.grid_n
460
461		maze_connection_lists: np.ndarray = np.empty(
462			(n_mazes, 2, grid_n, grid_n),
463			dtype=np.bool_,
464		)
465		# maze_endpoints: np.ndarray = np.empty((n_mazes, 2, 2), dtype=np.int8)
466		maze_solution_lengths: np.ndarray = np.empty((n_mazes,), dtype=np.int32)
467		maze_solutions: np.ndarray = np.empty(
468			(n_mazes, max_solution_len, 2),
469			dtype=np.int8,
470		)
471
472		for idx, maze in enumerate(filtered_meta.mazes):
473			maze_connection_lists[idx] = maze.connection_list
474			# maze_endpoints[idx] = np.array([maze.start_pos, maze.end_pos])
475			maze_solution_lengths[idx] = maze.solution.shape[0]
476			maze_solutions[idx, : maze.solution.shape[0]] = maze.solution
477
478		return {
479			_FORMAT_KEY: "MazeDataset:minimal",
480			"cfg": json_serialize(filtered_meta.cfg),
481			"fname": filtered_meta.cfg.to_fname(),
482			"generation_metadata_collected": json_serialize(
483				filtered_meta.generation_metadata_collected,
484			),
485			"maze_connection_lists": maze_connection_lists,  # type: ignore[dict-item]
486			# "maze_endpoints": maze_endpoints,
487			"maze_solution_lengths": maze_solution_lengths,  # type: ignore[dict-item]
488			"maze_solutions": maze_solutions,  # type: ignore[dict-item]
489		}
490
491	def _serialize_minimal_soln_cat(self: "MazeDataset") -> JSONdict:
492		"alternate serialization where metadata is collected, and mazes and their solutions are stored in concatenated form"
493		filtered_meta: MazeDataset
494		if self.generation_metadata_collected is None:
495			filtered_meta = self.filter_by.collect_generation_meta()
496		else:
497			filtered_meta = self
498
499		maze_solution_lengths: np.ndarray = np.array(
500			[m.solution.shape[0] for m in filtered_meta.mazes],
501			dtype=np.int32,
502		)
503		n_mazes: int = len(filtered_meta.mazes)
504		grid_n: int = filtered_meta.cfg.grid_n
505		total_solution_len: int = np.sum(maze_solution_lengths)
506
507		maze_connection_lists: np.ndarray = np.empty(
508			(n_mazes, 2, grid_n, grid_n),
509			dtype=np.bool_,
510		)
511		maze_endpoints: np.ndarray = np.empty((n_mazes, 2, 2), dtype=np.int8)
512		maze_solutions_concat: np.ndarray = np.empty(
513			(total_solution_len, 2),
514			dtype=np.int8,
515		)
516
517		solutions_running_idx: int = 0
518		for idx, maze in enumerate(filtered_meta.mazes):
519			maze_connection_lists[idx] = maze.connection_list
520			maze_endpoints[idx] = np.array([maze.start_pos, maze.end_pos])
521			soln_len: int = maze.solution.shape[0]
522			maze_solution_lengths[idx] = soln_len
523			maze_solutions_concat[
524				solutions_running_idx : solutions_running_idx + soln_len
525			] = maze.solution
526			solutions_running_idx += soln_len
527
528		return {
529			_FORMAT_KEY: "MazeDataset:minimal_soln_cat",
530			"cfg": json_serialize(filtered_meta.cfg),
531			"fname": filtered_meta.cfg.to_fname(),
532			"generation_metadata_collected": json_serialize(
533				filtered_meta.generation_metadata_collected,
534			),
535			"maze_connection_lists": maze_connection_lists,  # type: ignore[dict-item]
536			"maze_endpoints": maze_endpoints,  # type: ignore[dict-item]
537			"maze_solution_lengths": maze_solution_lengths,  # type: ignore[dict-item]
538			"maze_solutions_concat": maze_solutions_concat,  # type: ignore[dict-item]
539		}
540
541	def update_self_config(self) -> None:
542		"""update the config to match the current state of the dataset (number of mazes, such as after filtering)"""
543		if self.cfg.n_mazes != len(self.mazes):
544			warnings.warn(
545				f"updating config n_mazes from {self.cfg.n_mazes} to {len(self.mazes)}",
546			)
547			self.cfg.n_mazes = len(self.mazes)
548
549	def custom_maze_filter(
550		self,
551		method: typing.Callable[[SolvedMaze], bool],
552		**kwargs,
553	) -> "MazeDataset":
554		"""filter the dataset using a custom method"""
555		output: MazeDataset = MazeDataset(
556			cfg=copy.deepcopy(self.cfg),
557			mazes=[m for m in self.mazes if method(m, **kwargs)],
558		)
559		output.cfg.applied_filters.append(
560			{
561				"name": f"__custom__:{method.__name__}",
562				"kwargs": kwargs,
563			},
564		)
565		output.update_self_config()
566		return output

257@serializable_dataclass(kw_only=True, methods_no_override=["serialize"])
258class MazeDatasetConfig(MazeDatasetConfig_base):  # type: ignore[misc]
259	"""config object which is passed to `MazeDataset.from_config` to generate or load a dataset
260
261	# Parameters:
262	- `name : str`
263		name of the dataset -- this can be anything, but should be filesystem safe since we use it in the `fname`
264	- `grid_n : int`
265		grid size of the maze (number of rows/columns)
266	- `n_mazes : int`
267		number of mazes to request. For some combinations of `endpoint_kwargs` and `maze_ctor`, not all mazes might successfully generate.
268		see `EndpointKwargsType` for more details.
269	- `maze_ctor : Callable`
270		maze generator function. This should be a function that takes a grid size and returns a maze.
271		This will usually be one of the functions in `LatticeMazeGenerators`.
272	- `maze_ctor_kwargs : dict`
273		keyword arguments to pass to the maze generator function. Specific to the `maze_ctor` you are using.
274	- `endpoint_kwargs : EndpointKwargsType`
275		keyword arguments passed to `LatticeMaze.generate_random_path()`. see `EndpointKwargsType` for more info.
276	- `applied_filters : list[dict]`
277		list of filters that have been applied to the dataset. We recommend applying filters to datasets directly,
278		but these are stored with the config in case you want to re-generate the dataset with the same filters.
279
280	"""
281
282	@property
283	def config_version(self) -> str:
284		"""return the version of the config. added in maze_dataset v1.3.0, previous versions had no dataset config"""
285		return "1.0"
286
287	@property
288	def versions(self) -> dict:
289		"""return the versions of the config and the maze_dataset"""
290		return dict(
291			config=self.config_version,
292			maze_dataset=importlib.metadata.version("maze_dataset"),
293		)
294
295	def serialize(self) -> dict:
296		"serialize the MazeDatasetConfig with all fields and fname"
297		return {
298			**self._serialize_base(
299				applied_filters__skip__collect_generation_meta=False
300			),
301			"fname": self.to_fname(),
302			"versions": self.versions,
303		}
304
305	def summary(self) -> dict:
306		"""return a summary of the config"""
307		# do we run this to make sure it doesn't error?
308		super_summary: dict = super().summary()
309		assert super_summary
310		self_ser: dict = self.serialize()
311		return dict(
312			name=self.name,
313			fname=self.to_fname(),
314			sdc_hash=self.stable_hash_cfg(),
315			seed=self.seed,
316			seq_len_min=self.seq_len_min,
317			seq_len_max=self.seq_len_max,
318			applied_filters=self.applied_filters,
319			grid_n=self_ser["grid_n"],
320			n_mazes=self_ser["n_mazes"],
321			maze_ctor_name=self_ser["maze_ctor"]["__name__"],
322			maze_ctor_kwargs=self_ser["maze_ctor_kwargs"],
323			endpoint_kwargs=self_ser["endpoint_kwargs"],
324		)
325
326	def _to_ps_array(self) -> _PercolationSuccessArray:
327		"""Convert this config to a [p, grid_n, deadends, endpoints_not_equal, generator_func] vector.
328
329		used in predicting the success rate
330		"""
331		try:
332			assert self.maze_ctor.__name__ in _GENERATORS_PERCOLATED, (
333				f"generator not supported, must be a percolation generator\n{self.maze_ctor.__name__ = }, {_GENERATORS_PERCOLATED = }"
334			)
335			assert "p" in self.maze_ctor_kwargs, (
336				f"maze_ctor_kwargs must have a 'p' (percolation value) key: {self.maze_ctor_kwargs = }"
337			)
338			assert not self.endpoint_kwargs.get("except_on_no_valid_endpoint", True), (
339				f"except_on_no_valid_endpoint must be False, or else if any maze fails to generate, the whole dataset will fail: {self.endpoint_kwargs = }"
340			)
341		except AssertionError as e:
342			err_msg: str = f"invalid config for percolation success prediction: {self.summary() = }"
343			raise NoPercolationInConfigError(
344				err_msg,
345			) from e
346
347		endpoints_unique_flag: int = int(
348			# we are pretty sure it will be an int or bool here
349			self.endpoint_kwargs.get("endpoints_not_equal", True),  # type: ignore[arg-type]
350		)
351
352		# adjustment for bknutson0
353		if not (
354			self.endpoint_kwargs.get("deadend_start", False)
355			and self.endpoint_kwargs.get("deadend_end", False)
356		):
357			# we didnt train on this, but if either endpoint is not required to be in a dead end
358			# then  requiring the endpoints to be unique does not really affect the success rate
359			# (except for very small percolation values, pure percolation generation)
360			endpoints_unique_flag = 0
361
362		return np.array(
363			[
364				float(self.maze_ctor_kwargs["p"]),
365				float(self.grid_n),
366				float(
367					int(
368						self.endpoint_kwargs.get("deadend_start", False)  # type: ignore[arg-type]
369						or self.endpoint_kwargs.get("deadend_end", False),
370					),
371				),
372				float(endpoints_unique_flag),
373				float(_GENERATORS_PERCOLATED.index(self.maze_ctor.__name__)),
374			],
375			dtype=np.float64,
376		)
377
378	@classmethod
379	def _from_ps_array(
380		cls,
381		arr: _PercolationSuccessArray,
382		name: str = "predict",
383		n_mazes: int = 100,
384		**kwargs,
385	) -> "MazeDatasetConfig":
386		"""Reconstruct a config from an array [p, grid_n, deadends, endpoints_not_equal, generator_func] and other config parameters.
387
388		# Returns:
389		- `MazeDatasetConfig`
390			Config corresponding to `arr`
391		"""
392		return cls(
393			name=name,
394			grid_n=int(arr[1]),
395			n_mazes=n_mazes,
396			maze_ctor=GENERATORS_MAP[_GENERATORS_PERCOLATED[int(arr[4])]],
397			maze_ctor_kwargs={"p": float(arr[0])},
398			endpoint_kwargs=dict(
399				deadend_start=bool(arr[2]),
400				deadend_end=bool(arr[2]),
401				endpoints_not_equal=bool(arr[3]),
402				except_on_no_valid_endpoint=False,
403			),
404			**kwargs,
405		)
406
407	def success_fraction_estimate(
408		self,
409		except_if_all_success_expected: bool = False,
410	) -> float:
411		"""Estimate the success fraction of this config.
412
413		only valid when the generator is a percolation generator,
414		and endpoints are enforced to be dead ends
415
416		more information on where this comes from can be found in
417		- `cfg_success_predict_fn()` from `maze_dataset.dataset.success_predict_math`
418		- `estimate_dataset_fractions.ipynb`
419		- `maze_dataset.benchmarks.sweep_fit`
420
421		# Parameters:
422		- `except_if_all_success_expected : bool`
423			if `True`, don't raise an error if the success fraction is below the threshold.
424			will always return `1.0` if the config is not expected to fail
425
426		# Returns:
427		- `float`
428			estimated success fraction
429
430		# Raises:
431		- `NoPercolationInConfigError` : if the config is not expected to fail, and `except_if_all_success_expected` is `False`
432		"""
433		try:
434			return cfg_success_predict_fn(self)
435
436		except NoPercolationInConfigError as e:
437			if except_if_all_success_expected:
438				raise e  # noqa: TRY201
439			return 1.0
440
441	def success_fraction_compensate(
442		self,
443		safety_margin: float = 1.2,
444		except_if_all_success_expected: bool = False,
445		epsilon: float = 1e-2,
446	) -> "MazeDatasetConfig":
447		"""return a new `MazeDatasetConfig` like this one with `n_mazes` adjusted to compensate for the success fraction
448
449		calls `MazeDatasetConfig.success_fraction_estimate()` to get the success fraction, and then
450		computes the new number of mazes as `n_mazes = n_mazes * safety_margin / success_fraction + 1`
451
452		more information on where this comes from can be found in
453		- `cfg_success_predict_fn()` from `maze_dataset.dataset.success_predict_math`
454		- `estimate_dataset_fractions.ipynb`
455		- `maze_dataset.benchmarks.sweep_fit`
456
457		# Parameters:
458		- `safety_margin : float`
459			safety margin to apply to the success fraction estimate
460			(defaults to `1.2`, or 20% more mazes than estimated)
461		- `except_if_all_success_expected : bool`
462			if `True`, don't raise an error if the success fraction is below the threshold.
463			this is passed to `MazeDatasetConfig.success_fraction_estimate`.
464			if your config isn't expected to fail, passing this might mean you generate more mazes than needed
465			since `safety_margin` is still applied.
466			(defaults to `False`)
467		- `epsilon : float`
468			raise `SuccessChanceTooSmallError` if the success fraction is below this threshold
469			(defaults to `1e-2`)
470
471		# Returns:
472		- `MazeDatasetConfig`
473			new config with adjusted `n_mazes`
474
475		# Raises:
476		- `SuccessChanceTooSmallError` : if the computed success fraction is below `epsilon`
477		"""
478		# compute and check the success fraction
479		success_fraction: float = self.success_fraction_estimate(
480			except_if_all_success_expected=except_if_all_success_expected,
481		)
482		if success_fraction < epsilon:
483			err_msg: str = (
484				f"{success_fraction = } is below the threshold of {epsilon = }"
485			)
486			raise SuccessChanceTooSmallError(
487				err_msg,
488			)
489
490		# compute the new number of mazes
491		n_mazes: int = self.n_mazes
492		new_n_mazes: int = int((n_mazes * safety_margin) / success_fraction) + 1
493
494		# put it in a new config and return
495		cfg_dict: dict = self.serialize()
496		cfg_dict["n_mazes"] = new_n_mazes
497		return MazeDatasetConfig.load(cfg_dict)