src/experiments.rs

changeset 12
a0d61ac3719e
parent 3
c3a4f4bb87f7
equal deleted inserted replaced
11:c4f923017349 12:a0d61ac3719e
7 process_error, Differentiable, HasProx, PythonMapping, PythonPlotFactory, 7 process_error, Differentiable, HasProx, PythonMapping, PythonPlotFactory,
8 }; 8 };
9 use alg_tools::direct_product::Pair; 9 use alg_tools::direct_product::Pair;
10 use alg_tools::error::{DynError, DynResult}; 10 use alg_tools::error::{DynError, DynResult};
11 use alg_tools::loc::Loc; 11 use alg_tools::loc::Loc;
12 use alg_tools::sets::Cube;
12 use anyhow::bail; 13 use anyhow::bail;
13 use clap::Parser; 14 use clap::Parser;
14 use log::debug; 15 use log::debug;
16 use pointsource_algs::domain::{AdmissibleDomain, FullSpace, UniformRandom};
15 use pointsource_algs::measures::DiscreteMeasure; 17 use pointsource_algs::measures::DiscreteMeasure;
16 use pointsource_algs::prox_penalty::RadonSquared; 18 use pointsource_algs::prox_penalty::RadonSquared;
17 use pointsource_algs::regularisation::{NonnegRadonRegTerm, RadonRegTerm, Regularisation}; 19 use pointsource_algs::regularisation::{NonnegRadonRegTerm, RadonRegTerm, Regularisation};
18 use pointsource_algs::run::{ 20 use pointsource_algs::run::{
19 run_fb, run_fb_pair, AlgorithmConfig, DefaultAlgorithm, Named, RunError::NotImplemented, 21 run_fb_dom, run_fb_pair_dom, DefaultAlgorithm, RunError::NotImplemented, RunnableExperiment,
20 RunnableExperiment, RunnableExperimentExtras, 22 RunnableExperimentExtras,
21 }; 23 };
22 use pointsource_algs::{AlgorithmOverrides, CommandLineArgs, ExperimentSetup}; 24 use pointsource_algs::{AlgorithmOverrides, CommandLineArgs, ExperimentSetup};
23 use pyo3::ffi::c_str; 25 use pyo3::ffi::c_str;
24 use pyo3::prelude::*; 26 use pyo3::prelude::*;
25 use pyo3::types::{PyDict, PyFunction}; 27 use pyo3::types::{PyDict, PyFunction};
28 use rand::prelude::{SeedableRng, StdRng};
26 use serde::{Deserialize, Serialize}; 29 use serde::{Deserialize, Serialize};
27 use serde_json; 30 use serde_json;
28 use serde_with::skip_serializing_none; 31 use serde_with::skip_serializing_none;
29 use std::collections::HashMap; 32 use std::collections::HashMap;
30 use std::ffi::CString; 33 use std::ffi::CString;
262 μinit: Option<DiscreteMeasure<Loc<2, f64>, f64>>, 265 μinit: Option<DiscreteMeasure<Loc<2, f64>, f64>>,
263 266
264 /// Plotter 267 /// Plotter
265 #[pyo3(set, get)] 268 #[pyo3(set, get)]
266 plot_factory: Option<Py<PyAny>>, 269 plot_factory: Option<Py<PyAny>>,
270
271 /// Domain
272 #[pyo3(set, get)]
273 domain: Option<Py<PyAny>>,
274
275 /// Randon number generator seed
276 #[pyo3(set, get)]
277 seed: Option<u64>,
267 } 278 }
268 279
269 #[pymethods] 280 #[pymethods]
270 impl Problem { 281 impl Problem {
271 #[new] 282 #[new]
272 #[pyo3(signature = (dataterm, regterm, auxterm = None, auxinit = None, μinit = None, plot_factory=None))] 283 #[pyo3(signature = (dataterm, regterm, auxterm = None, auxinit = None, μinit = None, plot_factory=None, domain = None, seed = None))]
273 fn new( 284 fn new(
274 dataterm: Py<PyAny>, 285 dataterm: Py<PyAny>,
275 regterm: RegTermPy, 286 regterm: RegTermPy,
276 auxterm: Option<Py<PyAny>>, 287 auxterm: Option<Py<PyAny>>,
277 auxinit: Option<Py<PyAny>>, 288 auxinit: Option<Py<PyAny>>,
278 μinit: Option<DiscreteMeasure<Loc<2, f64>, f64>>, 289 μinit: Option<DiscreteMeasure<Loc<2, f64>, f64>>,
279 plot_factory: Option<Py<PyAny>>, 290 plot_factory: Option<Py<PyAny>>,
291 domain: Option<Py<PyAny>>,
292 seed: Option<u64>,
280 ) -> Self { 293 ) -> Self {
281 Self { dataterm, regterm, auxterm, auxinit, μinit, plot_factory } 294 Self { dataterm, regterm, auxterm, auxinit, μinit, plot_factory, seed, domain }
282 } 295 }
283 296
284 #[getter] 297 #[getter]
285 fn get_dataterm<'py>(&self, py: Python<'py>) -> Bound<'py, PyAny> { 298 fn get_dataterm<'py>(&self, py: Python<'py>) -> Bound<'py, PyAny> {
286 self.dataterm.bind(py).clone() 299 self.dataterm.bind(py).clone()
287 } 300 }
288 } 301 }
289 302
303 impl PythonExperiment {
304 fn do_runall_phase2<D: AdmissibleDomain<Loc<2, f64>>>(
305 &self,
306 cli: &CommandLineArgs,
307 cli_algorithm_overrides: &AlgorithmOverrides<f64>,
308 prefix: String,
309 problem: PyRef<Problem>,
310 domain: &mut D,
311 py: Python,
312 ) -> DynError {
313 // Set up algorithms
314 let algorithms = self.collect_algs(
315 cli,
316 cli_algorithm_overrides,
317 |alg| self.algorithm_overrides.get(&alg).cloned(),
318 Some(&problem.regterm.into()),
319 &[DefaultAlgorithm::RadonFB, DefaultAlgorithm::RadonSlidingFB],
320 )?;
321
322 let save_extra = |_, ()| Ok(());
323 let μ0 = problem.μinit.clone();
324
325 match (problem.regterm, &problem.auxterm, &problem.auxinit) {
326 (regterm, None, None) => {
327 debug!("… Extracting data term.");
328 let dataterm: exp_f64_2::DataTerm<'_> = problem.dataterm.extract(py)?;
329 debug!("… Extracting plotter.");
330 let plot_factory: exp_f64_2::PlotFactory<'_> =
331 if let Some(ref p) = problem.plot_factory {
332 p.extract(py)?
333 } else {
334 PythonPlotFactory::dummy()
335 };
336 let make_plotter = |prefix| plot_factory.produce(prefix);
337
338 debug!("Running.");
339 self.do_runall(
340 &prefix,
341 cli,
342 algorithms,
343 make_plotter,
344 save_extra,
345 μ0,
346 |p| match regterm {
347 RegTermPy::Radon(α) => run_fb_dom(
348 &dataterm,
349 &RadonRegTerm(α),
350 &RadonSquared,
351 p,
352 domain,
353 |_| Err(NotImplemented.into()),
354 )
355 .map(|μ| (μ, ())),
356 RegTermPy::NonnegRadon(α) => run_fb_dom(
357 &dataterm,
358 &NonnegRadonRegTerm(α),
359 &RadonSquared,
360 p,
361 domain,
362 |_| Err(NotImplemented.into()),
363 )
364 .map(|μ| (μ, ())),
365 },
366 )
367 }
368 (regterm, Some(auxterm), Some(z_)) => {
369 debug!("… Extracting auxiliary variable.");
370 let z0: PyResult<exp_f64_2_auxvar::AuxVar<'_>> = z_.extract(py);
371
372 match z0 {
373 Ok(z) => {
374 debug!("… Extracting data term.");
375 let dataterm: exp_f64_2_auxvar::DataTerm<'_> =
376 problem.dataterm.extract(py)?;
377 debug!("… Extracting plotter.");
378 let plot_factory: exp_f64_2::PlotFactory<'_> =
379 if let Some(ref p) = problem.plot_factory {
380 p.extract(py)?
381 } else {
382 PythonPlotFactory::dummy()
383 };
384 let make_plotter = |prefix| plot_factory.produce(prefix);
385 debug!("… Extracting auxiliary term.");
386 let auxterm: exp_f64_2_auxvar::AuxTerm<'_> = auxterm.extract(py)?;
387
388 debug!("Running.");
389 self.do_runall(
390 &prefix,
391 cli,
392 algorithms,
393 make_plotter,
394 save_extra,
395 (μ0, z),
396 |p| match regterm {
397 RegTermPy::Radon(α) => run_fb_pair_dom(
398 &dataterm,
399 &RadonRegTerm(α),
400 &RadonSquared,
401 &auxterm,
402 p,
403 domain,
404 |_| Err(NotImplemented.into()),
405 )
406 .map(|Pair(μ, _)| (μ, ())),
407 RegTermPy::NonnegRadon(α) => run_fb_pair_dom(
408 &dataterm,
409 &NonnegRadonRegTerm(α),
410 &RadonSquared,
411 &auxterm,
412 p,
413 domain,
414 |_| Err(NotImplemented.into()),
415 )
416 .map(|Pair(μ, _)| (μ, ())),
417 },
418 )
419 }
420 Err(_) => {
421 let z: exp_f64_2_auxvar_scalar::AuxVar<'_> = z_.extract(py)?;
422 debug!("… Extracting data term.");
423 let dataterm: exp_f64_2_auxvar_scalar::DataTerm<'_> =
424 problem.dataterm.extract(py)?;
425 debug!("… Extracting plotter.");
426 let plot_factory: exp_f64_2::PlotFactory<'_> =
427 if let Some(ref p) = problem.plot_factory {
428 p.extract(py)?
429 } else {
430 PythonPlotFactory::dummy()
431 };
432 let make_plotter = |prefix| plot_factory.produce(prefix);
433 debug!("… Extracting auxiliary term.");
434 let auxterm: exp_f64_2_auxvar_scalar::AuxTerm<'_> = auxterm.extract(py)?;
435
436 debug!("Running.");
437 self.do_runall(
438 &prefix,
439 cli,
440 algorithms,
441 make_plotter,
442 save_extra,
443 (μ0, z),
444 |p| match regterm {
445 RegTermPy::Radon(α) => run_fb_pair_dom(
446 &dataterm,
447 &RadonRegTerm(α),
448 &RadonSquared,
449 &auxterm,
450 p,
451 domain,
452 |_| Err(NotImplemented.into()),
453 )
454 .map(|Pair(μ, _)| (μ, ())),
455 RegTermPy::NonnegRadon(α) => run_fb_pair_dom(
456 &dataterm,
457 &NonnegRadonRegTerm(α),
458 &RadonSquared,
459 &auxterm,
460 p,
461 domain,
462 |_| Err(NotImplemented.into()),
463 )
464 .map(|Pair(μ, _)| (μ, ())),
465 },
466 )
467 }
468 }
469 }
470 (_, _, _) => {
471 bail!("Errors in problem {} setup: auxiliary term or auxiliary variable initialisation given without the other", self.name);
472 }
473 }
474 }
475 }
476
290 impl RunnableExperiment<f64> for PythonExperiment { 477 impl RunnableExperiment<f64> for PythonExperiment {
291 fn name(&self) -> &str { 478 fn name(&self) -> &str {
292 &self.name 479 &self.name
293 } 480 }
294 481
295 fn runall( 482 fn runall(
296 &self, 483 &self,
297 cli: &CommandLineArgs, 484 cli: &CommandLineArgs,
298 algs: Option<Vec<Named<AlgorithmConfig<f64>>>>, 485 cli_algorithm_overrides: &AlgorithmOverrides<f64>,
299 ) -> DynError { 486 ) -> DynError {
300 // Set up algorithms
301 let algorithms = algs.unwrap_or_else(|| vec![DefaultAlgorithm::FB.get_named()]);
302
303 debug!( 487 debug!(
304 "Entered PythonExperiment::runall for experimen {}. Attaching Python.", 488 "Entered PythonExperiment::runall for experimen {}. Attaching Python.",
305 self.name 489 self.name
306 ); 490 );
307 491
308 Python::attach(|py| { 492 Python::attach(|py| {
309 debug!("Calling Python-side experiment initialisation."); 493 debug!("Calling Python-side experiment initialisation.");
310 494
311 let prefix = self.start(cli)?; 495 let prefix = self.start(cli)?;
312 496
497 // Get the problem description
313 let problem: PyRef<Problem> = process_error( 498 let problem: PyRef<Problem> = process_error(
314 "PythonExperiment::runall", 499 "PythonExperiment::runall",
315 py, 500 py,
316 self.setup 501 self.setup
317 .call1(py, (&prefix,)) 502 .call1(py, (&prefix,))
318 .and_then(|r| r.extract(py).map_err(PyErr::from)), 503 .and_then(|r| r.extract(py).map_err(PyErr::from)),
319 )?; 504 )?;
320 505
321 let save_extra = |_, ()| Ok(()); 506 let rng = StdRng::seed_from_u64(problem.seed.unwrap_or(42));
322 let μ0 = problem.μinit.clone(); 507 if let Some(ref d) = problem.domain {
323 508 let mut domain: UniformRandom<Cube<2, f64>, _> =
324 match (problem.regterm, &problem.auxterm, &problem.auxinit) { 509 UniformRandom(d.extract::<[[f64; 2]; 2]>(py)?.into(), rng);
325 (regterm, None, None) => { 510 self.do_runall_phase2(
326 debug!("… Extracting data term."); 511 cli,
327 let dataterm: exp_f64_2::DataTerm<'_> = problem.dataterm.extract(py)?; 512 cli_algorithm_overrides,
328 debug!("… Extracting plotter."); 513 prefix,
329 let plot_factory: exp_f64_2::PlotFactory<'_> = 514 problem,
330 if let Some(ref p) = problem.plot_factory { 515 &mut domain,
331 p.extract(py)? 516 py,
332 } else { 517 )
333 PythonPlotFactory::dummy() 518 } else {
334 }; 519 self.do_runall_phase2(
335 let make_plotter = |prefix| plot_factory.produce(prefix); 520 cli,
336 521 cli_algorithm_overrides,
337 debug!("Running."); 522 prefix,
338 self.do_runall( 523 problem,
339 &prefix, 524 &mut FullSpace,
340 cli, 525 py,
341 algorithms, 526 )
342 make_plotter,
343 save_extra,
344 μ0,
345 |p| match regterm {
346 RegTermPy::Radon(α) => {
347 run_fb(&dataterm, &RadonRegTerm(α), &RadonSquared, p, |_| {
348 Err(NotImplemented.into())
349 })
350 .map(|μ| (μ, ()))
351 }
352 RegTermPy::NonnegRadon(α) => {
353 run_fb(&dataterm, &NonnegRadonRegTerm(α), &RadonSquared, p, |_| {
354 Err(NotImplemented.into())
355 })
356 .map(|μ| (μ, ()))
357 }
358 },
359 )
360 }
361 (regterm, Some(auxterm), Some(z_)) => {
362 debug!("… Extracting auxiliary variable.");
363 let z0: PyResult<exp_f64_2_auxvar::AuxVar<'_>> = z_.extract(py);
364 match z0 {
365 Ok(z) => {
366 debug!("… Extracting data term.");
367 let dataterm: exp_f64_2_auxvar::DataTerm<'_> =
368 problem.dataterm.extract(py)?;
369 debug!("… Extracting plotter.");
370 let plot_factory: exp_f64_2::PlotFactory<'_> =
371 if let Some(ref p) = problem.plot_factory {
372 p.extract(py)?
373 } else {
374 PythonPlotFactory::dummy()
375 };
376 let make_plotter = |prefix| plot_factory.produce(prefix);
377 debug!("… Extracting auxiliary term.");
378 let auxterm: exp_f64_2_auxvar::AuxTerm<'_> = auxterm.extract(py)?;
379
380 debug!("Running.");
381 self.do_runall(
382 &prefix,
383 cli,
384 algorithms,
385 make_plotter,
386 save_extra,
387 (μ0, z),
388 |p| match regterm {
389 RegTermPy::Radon(α) => run_fb_pair(
390 &dataterm,
391 &RadonRegTerm(α),
392 &RadonSquared,
393 &auxterm,
394 p,
395 |_| Err(NotImplemented.into()),
396 )
397 .map(|Pair(μ, _)| (μ, ())),
398 RegTermPy::NonnegRadon(α) => run_fb_pair(
399 &dataterm,
400 &NonnegRadonRegTerm(α),
401 &RadonSquared,
402 &auxterm,
403 p,
404 |_| Err(NotImplemented.into()),
405 )
406 .map(|Pair(μ, _)| (μ, ())),
407 },
408 )
409 }
410 Err(_) => {
411 let z: exp_f64_2_auxvar_scalar::AuxVar<'_> = z_.extract(py)?;
412 debug!("… Extracting data term.");
413 let dataterm: exp_f64_2_auxvar_scalar::DataTerm<'_> =
414 problem.dataterm.extract(py)?;
415 debug!("… Extracting plotter.");
416 let plot_factory: exp_f64_2::PlotFactory<'_> =
417 if let Some(ref p) = problem.plot_factory {
418 p.extract(py)?
419 } else {
420 PythonPlotFactory::dummy()
421 };
422 let make_plotter = |prefix| plot_factory.produce(prefix);
423 debug!("… Extracting auxiliary term.");
424 let auxterm: exp_f64_2_auxvar_scalar::AuxTerm<'_> =
425 auxterm.extract(py)?;
426
427 debug!("Running.");
428 self.do_runall(
429 &prefix,
430 cli,
431 algorithms,
432 make_plotter,
433 save_extra,
434 (μ0, z),
435 |p| match regterm {
436 RegTermPy::Radon(α) => run_fb_pair(
437 &dataterm,
438 &RadonRegTerm(α),
439 &RadonSquared,
440 &auxterm,
441 p,
442 |_| Err(NotImplemented.into()),
443 )
444 .map(|Pair(μ, _)| (μ, ())),
445 RegTermPy::NonnegRadon(α) => run_fb_pair(
446 &dataterm,
447 &NonnegRadonRegTerm(α),
448 &RadonSquared,
449 &auxterm,
450 p,
451 |_| Err(NotImplemented.into()),
452 )
453 .map(|Pair(μ, _)| (μ, ())),
454 },
455 )
456 }
457 }
458 }
459 (_, _, _) => {
460 bail!("Errors in problem {} setup: auxiliary term or auxiliary variable initialisation given without the other", self.name);
461 }
462 } 527 }
463 })?; 528 })
464 529 }
465 // Should run the experiment here, going through all algorithms. 530 }
466 Ok(())
467 }
468
469 fn algorithm_overrides(&self, alg: DefaultAlgorithm) -> AlgorithmOverrides<f64> {
470 self.algorithm_overrides
471 .get(&alg)
472 .cloned()
473 .unwrap_or(Default::default())
474 }
475 }

mercurial